Final version of pycbc_pygrb_plot_snr_timeseries (#4955)

* Handling new option for pycbc_pygrb_plot_snr_timeseries jobs

* Avoid single use of variable

bin/pygrb/pycbc_pygrb_plot_snr_timeseries

@@ -48,30 +48,6 @@ __program__ = "pycbc_pygrb_plot_snr_timeseries"
 # =============================================================================
 # Functions
 # =============================================================================
-# Load trigger data
-def load_data(input_file, ifos, vetoes, rw_snr_threshold=None,
-              injections=False, slide_id=None):
-    """Load data from a trigger/injection file"""
-
-    trigs_or_injs = None
-    if input_file:
-        if injections:
-            logging.info("Loading injections...")
-            # This will eventually become load_injections
-            trigs_or_injs = \
-                ppu.load_triggers(input_file, ifos, vetoes,
-                                  rw_snr_threshold=rw_snr_threshold,
-                                  slide_id=slide_id)
-        else:
-            logging.info("Loading triggers...")
-            trigs_or_injs = \
-                ppu.load_triggers(input_file, ifos, vetoes,
-                                  rw_snr_threshold=rw_snr_threshold,
-                                  slide_id=slide_id)
-
-    return trigs_or_injs
-
-
 # Find start and end times of trigger/injecton data relative to a given time
 def get_start_end_times(data_time, central_time):
     """Determine padded start and end times of data relative to central_time"""
@@ -108,6 +84,8 @@ parser.add_argument("--trigger-time", type=float, default=0,
 parser.add_argument("-y", "--y-variable", default=None,
                     choices=['coherent', 'single', 'reweighted', 'null'],
                     help="Quantity to plot on the vertical axis.")
+parser.add_argument("--onsource", default=False, action="store_true",
+                    help="Include onsource data in the plot (CAUTION!)")
 ppu.pygrb_add_bestnr_cut_opt(parser)
 ppu.pygrb_add_slide_opts(parser)
 opts = parser.parse_args()
@@ -132,48 +110,83 @@ outdir = os.path.split(os.path.abspath(opts.output_file))[0]
 if not os.path.isdir(outdir):
     os.makedirs(outdir)
 
-# Extract IFOs and vetoes
-ifos, vetoes = ppu.extract_ifos_and_vetoes(trig_file, opts.veto_files,
-                                           opts.veto_category)
+# Extract IFOs
+ifos = ppu.extract_ifos(trig_file)
+
+# Generate time-slides dictionary
+slide_dict = ppu.load_time_slides(trig_file)
+
+# Generate segments dictionary
+segment_dict = ppu.load_segment_dict(trig_file)
+
+# Construct trials removing vetoed times
+trial_dict, total_trials = ppu.construct_trials(
+    opts.seg_files,
+    segment_dict,
+    ifos,
+    slide_dict,
+    opts.veto_file,
+    hide_onsource=(not opts.onsource)
+)
 
 # Load trigger and injections data: when plotting reweighted SNR, keep all
 # points to show the impact of the cut, otherwise remove points with
 # reweighted SNR below threshold
-if snr_type == 'reweighted':
-    trig_data = load_data(trig_file, ifos, vetoes,
+rw_snr_threshold = None if snr_type == 'reweighted' else opts.newsnr_threshold
+# When including the onsource, avoid printing to screen via logging the number
+# of triggers found
+data_tag = 'trigs' if opts.onsource else None
+trig_data = ppu.load_data(trig_file, ifos, data_tag=data_tag,
+                          rw_snr_threshold=rw_snr_threshold,
                           slide_id=opts.slide_id)
-    trig_data['network/reweighted_snr'] = \
-        reweightedsnr_cut(trig_data['network/reweighted_snr'],
-                          opts.newsnr_threshold)
-    inj_data = load_data(inj_file, ifos, vetoes, injections=True,
+inj_data = ppu.load_data(inj_file, ifos, data_tag='injs',
+                         rw_snr_threshold=rw_snr_threshold,
                          slide_id=0)
-    if inj_data is not None:
-        inj_data['network/reweighted_snr'] = \
-            reweightedsnr_cut(inj_data['network/reweighted_snr'],
-                              opts.newsnr_threshold)
-else:
-    trig_data = load_data(trig_file, ifos, vetoes,
-                          rw_snr_threshold=opts.newsnr_threshold,
-                          slide_id=opts.slide_id)
-    inj_data = load_data(inj_file, ifos, vetoes,
-                         rw_snr_threshold=opts.newsnr_threshold,
-                         injections=True, slide_id=0)
 
 # Specify HDF file keys for x quantity (time) and y quantity (SNR)
-if snr_type == 'single':
+if opts.y_variable == 'single':
     x_key = opts.ifo + '/end_time'
     y_key = opts.ifo + '/snr'
 else:
     x_key = 'network/end_time_gc'
-    y_key = 'network/' + snr_type + '_snr'
-
-# Obtain times
-trig_data_time = trig_data[x_key][:]
-inj_data_time = inj_data[x_key][:] if inj_file else None
-
-# Obtain SNRs
-trig_data_snr = trig_data[y_key][:]
-inj_data_snr = inj_data[y_key][:] if inj_file else None
+    y_key = 'network/' + opts.y_variable + '_snr'
+
+# Extract needed trigger properties and store them as dictionaries
+# Based on trial_dict: if vetoes were applied, trig_* are the veto survivors
+found_trigs = ppu.extract_trig_properties(
+    trial_dict,
+    trig_data,
+    slide_dict,
+    segment_dict,
+    [x_key, y_key]
+)
+
+# Gather injections found surviving vetoes
+found_after_vetoes, *_ = ppu.apply_vetoes_to_found_injs(
+    opts.found_missed_file,
+    inj_data,
+    ifos,
+    veto_file=opts.veto_file,
+    keys=[x_key, y_key]
+)
+
+# Obtain times and SNRs
+trig_data_time = numpy.concatenate([found_trigs[x_key][slide_id][:] for slide_id in slide_dict])
+trig_data_snr = numpy.concatenate([found_trigs[y_key][slide_id][:] for slide_id in slide_dict])
+inj_data_time = found_after_vetoes[x_key][:] if inj_file else None
+inj_data_snr = found_after_vetoes[y_key][:] if inj_file else None
+
+# Apply reweighted SNR threshold keeping downweighted points
+if snr_type == 'reweighted':
+    trig_data_snr = reweightedsnr_cut(
+        trig_data_snr,
+        opts.newsnr_threshold
+    )
+    if inj_file:
+        inj_data_snr = reweightedsnr_cut(
+            inj_data_snr,
+            opts.newsnr_threshold
+        )
 
 # Determine the central time (t=0) for the plot
 central_time = opts.trigger_time
@@ -191,7 +204,7 @@ logging.info("Plotting...")
 
 # Determine what goes on the vertical axis
 y_labels = {'coherent': "Coherent SNR",
-            'single': "%s SNR" % ifo,
+            'single': f"{ifo} SNR",
             'null': "Null SNR",
             'reweighted': "Reweighted SNR"}
 y_label = y_labels[snr_type]
@@ -205,11 +218,9 @@ if opts.plot_caption is None:
         opts.plot_caption += ("Red crosses: injections triggers.")
 
 # Single IFO SNR versus time plots
-xlims = [start, end]
-if opts.x_lims:
-    xlims = opts.x_lims
-    xlims = map(float, xlims.split(','))
+if not opts.x_lims:
+    opts.x_lims = f"{start},{end}"
 plu.pygrb_plotter([trig_data_time, trig_data_snr],
                   [inj_data_time, inj_data_snr],
-                  "Time since %.3f (s)" % (central_time), y_label,
+                  f"Time since {central_time:.3f} (s)", y_label,
                   opts, cmd=' '.join(sys.argv))

pycbc/workflow/grb_utils.py

@@ -483,6 +483,9 @@ def make_pygrb_plot(workflow, exec_name, out_dir,
         node.add_input_list_opt('--seg-files', seg_files)
     if veto_file:
         node.add_input_opt('--veto-file', veto_file)
+    # Option to show the onsource trial if this is a plot of all data
+    if exec_name == 'pygrb_plot_snr_timeseries' and 'alltimes' in tags:
+        node.add_opt('--onsource')
     if exec_name in ['pygrb_plot_injs_results',
                      'pygrb_plot_snr_timeseries']:
         trig_time = workflow.cp.get('workflow', 'trigger-time')