Fix 4491 (#4510)

bin/minifollowups/pycbc_plot_trigger_timeseries

@@ -18,7 +18,7 @@
 import argparse, logging, pycbc.version, pycbc.results, sys
 from pycbc.types import MultiDetOptionAction
 from pycbc.events import ranking
-from pycbc.io import HFile
+from pycbc.io import HFile, SingleDetTriggers
 import h5py
 import matplotlib; matplotlib.use('Agg'); import pylab
 import numpy
@@ -58,11 +58,14 @@ for ifo in args.single_trigger_files.keys():
     logging.info("Getting %s triggers", ifo)
     t = args.times[ifo]
 
-    data = HFile(args.single_trigger_files[ifo], 'r')
+    # Identify trigger idxs within window of trigger time
+    with HFile(args.single_trigger_files[ifo], 'r') as data:
+        idx, _ = data.select(lambda endtime: abs(endtime - t) < args.window,
+                             f'{ifo}/end_time', return_indices=True)
+        data_mask = numpy.zeros(len(data[f'{ifo}/snr']), dtype=bool)
+        data_mask[idx] = True
 
-    # Identify trigger indices within window
-    keep = abs(data[ifo + '/end_time'][:] - t) < args.window
-    if not any(keep):
+    if not len(idx):
         # No triggers in this window, add to the legend and continue
         # Make sure it isnt on the plot
         pylab.scatter(-2 * args.window, 0,
@@ -71,16 +74,22 @@ for ifo in args.single_trigger_files.keys():
                       label=ifo)
         continue
 
-    any_data = True
-    logging.info("Keeping %d triggers in the window", sum(keep))
-
+    # Not using bank file, or veto file, so lots of "None"s here.
+    trigs = SingleDetTriggers(
+        args.single_trigger_files[ifo],
+        None,
+        None,
+        None,
+        None,
+        ifo,
+        premask=data_mask
+    )
 
-    data_dict = {k: data[ifo][k][keep] for k in data[ifo]
-                 if isinstance(data[ifo][k], h5py.Dataset)
-                 and data[ifo][k].size == keep.size}
+    any_data = True
+    logging.info("Keeping %d triggers in the window", len(idx))
 
     logging.info("Getting %s", args.plot_type)
-    rank = ranking.get_sngls_ranking_from_trigs(data_dict, args.plot_type)
+    rank = ranking.get_sngls_ranking_from_trigs(trigs, args.plot_type)
 
     if all(rank == 1):
         # This is the default value to say "downranked beyond consideration"
@@ -91,21 +100,23 @@ for ifo in args.single_trigger_files.keys():
                       label=ifo)
         continue
 
-    pylab.scatter(data_dict['end_time'] - t, rank,
+    pylab.scatter(trigs['end_time'] - t, rank,
                   color=pycbc.results.ifo_color(ifo), marker='x',
                   label=ifo)
 
     # Minimum rank which hasn't been set to the default of 1
     min_rank = min(min_rank, rank[rank > 1].min())
 
     if args.special_trigger_ids:
-        special_idx = numpy.flatnonzero(keep) == args.special_trigger_ids[ifo]
-        if not any(special_idx):
-            logging.info("IDX %d not in kept list", args.special_trigger_ids[ifo])
+        special_idx = args.special_trigger_ids[ifo]
+        if special_idx not in idx:
+            logging.info("IDX %d not in kept list",
+                         args.special_trigger_ids[ifo])
             continue
+        special_red_idx = numpy.where(idx == special_idx)[0]
 
-        pylab.scatter((data_dict['end_time'] - t)[special_idx],
-                      rank[special_idx], marker='*', s=50, color='yellow')
+        pylab.scatter(trigs.trigs_f[f'{ifo}/end_time'][special_idx] - t,
+                      rank[special_red_idx], marker='*', s=50, color='yellow')
 
 if args.log_y_axis and any_data:
     pylab.yscale('log')