Version of pycbc_pygrb_plot_stats_distribution vetoes

bin/pygrb/pycbc_pygrb_plot_stats_distribution

@@ -54,7 +54,9 @@ parser.add_argument("-x", "--x-variable", required=True,
                     choices=["bestnr", "snr", "snruncut"],
                     help="Quantity to plot on the horizontal axis.")
 ppu.pygrb_add_bestnr_cut_opt(parser)
+ppu.pygrb_add_slide_opts(parser)
 opts = parser.parse_args()
+ppu.slide_opts_helper(opts)
 
 init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s")
 
@@ -68,31 +70,45 @@ 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)
 
-# Load triggers, time-slides, and segment dictionary
-logging.info("Loading triggers.")
-trigs = ppu.load_triggers(trig_file, ifos, vetoes,
-                          rw_snr_threshold=opts.newsnr_threshold)
-logging.info("%d triggers loaded.", len(trigs['network/event_id']))
-logging.info("Loading timeslides.")
+# Generate time-slides dictionary
 slide_dict = ppu.load_time_slides(trig_file)
-logging.info("Loading segments.")
-segment_dict = ppu.load_segment_dict(trig_file)
 
-# Construct trials
-logging.info("Constructing trials.")
-trial_dict = ppu.construct_trials(opts.seg_files, segment_dict,
-                                  ifos, slide_dict, vetoes)
-total_trials = sum([len(trial_dict[slide_id]) for slide_id in slide_dict])
-logging.info("%d trials generated.", total_trials)
+# We will be looping over slide_dict below so here we reduce it if possible
+if opts.slide_id is not None:
+    for key in list(slide_dict.keys()):
+        if key != opts.slide_id:
+            slide_dict.pop(key, None)
+
+# Generate segments dictionary
+segment_dict = ppu.load_segment_dict(trig_file)
 
-# Extract basic trigger properties and store as dictionaries
-trig_time, trig_snr, trig_bestnr = \
-    ppu.extract_basic_trig_properties(trial_dict, trigs, slide_dict,
-                                      segment_dict, opts)
+# Construct trials removing vetoed times
+trial_dict, total_trials = ppu.construct_trials(
+    opts.seg_files,
+    segment_dict,
+    ifos,
+    slide_dict,
+    opts.veto_file
+)
+
+# Load triggers (apply reweighted SNR cut, not vetoes)
+all_off_trigs = ppu.load_data(trig_file, ifos, data_tag='trigs',
+                              rw_snr_threshold=opts.newsnr_threshold,
+                              slide_id=opts.slide_id)
+
+# Extract needed trigger properties and store them as dictionaries
+# Based on trial_dict: if vetoes were applied, trig_* are the veto survivors
+keys = ['network/end_time_gc', 'network/coherent_snr', 'network/reweighted_snr']
+trig_data = ppu.extract_trig_properties(
+    trial_dict,
+    all_off_trigs,
+    slide_dict,
+    segment_dict,
+    keys
+)
 
 # Calculate SNR and BestNR values and maxima
 time_veto_max_snr = {}
@@ -107,28 +123,28 @@ for slide_id in slide_dict:
 
 for slide_id in slide_dict:
     for j, trial in enumerate(trial_dict[slide_id]):
-        trial_cut = (trial[0] <= trig_time[slide_id])\
-                          & (trig_time[slide_id] < trial[1])
+        trial_cut = (trial[0] <= trig_data[keys[0]][slide_id][:])\
+                          & (trig_data[keys[0]][slide_id][:] < trial[1])
         if not trial_cut.any():
             continue
         # Max SNR
         time_veto_max_snr[slide_id][j] = \
-            max(trig_snr[slide_id][trial_cut])
+            max(trig_data[keys[1]][slide_id][trial_cut])
         # Max BestNR
         time_veto_max_bestnr[slide_id][j] = \
-            max(trig_bestnr[slide_id][trial_cut])
+            max(trig_data[keys[2]][slide_id][trial_cut])
         # Max SNR for triggers passing SBVs
-        sbv_cut = trig_bestnr[slide_id] != 0
+        sbv_cut = trig_data[keys[2]][slide_id][:] != 0
         if not (trial_cut & sbv_cut).any():
             continue
         time_veto_max_snr_uncut[slide_id][j] =\
-            max(trig_snr[slide_id][trial_cut & sbv_cut])
+            max(trig_data[keys[1]][slide_id][trial_cut & sbv_cut])
 
 # This is the data that will be plotted
 full_time_veto_max_snr = ppu.sort_stat(time_veto_max_snr)
 full_time_veto_max_snr_uncut = ppu.sort_stat(time_veto_max_snr_uncut)
 _, _, full_time_veto_max_bestnr = \
-    ppu.max_median_stat(slide_dict, time_veto_max_bestnr, trig_bestnr,
+    ppu.max_median_stat(slide_dict, time_veto_max_bestnr, trig_data[keys[2]],
                         total_trials)
 # The 0.'s here force the histograms to start at (0, 1) if no trial
 # returned a no-event (i.e., BestNR = 0)
@@ -203,8 +219,7 @@ ax.fill_between(x + dx/2, lower, upper, alpha=0.3, facecolor='y', step='mid')
 ax.set_xlim((0.9 * min_stat, 1.1 * max_stat))
 ax.set_ylim((0.6/N, ymax))
 plot_title = "Cumulative distribution of background triggers"
-plot_caption = "Background cumulative distribution of the %s " % \
-    x_label_dict[stat]
+plot_caption = f"Background cumulative distribution of the {x_label_dict[stat]}"
 fig_path = opts.output_file
 save_fig_with_metadata(fig, fig_path, cmd=' '.join(sys.argv),
                        title=plot_title, caption=plot_caption)