Change to frame finding in collated DQ script (#4910)

* Fixes to collated dq script to allow using frames from llcache

* Use glob to gather frames

* update supervision script

* Remove debug statements

* Fixed problem with argument type

* Make gwdatafind usable again

bin/live/pycbc_live_collated_dq_trigger_rates

@@ -20,63 +20,124 @@ for a day of PyCBC Live triggers.
 import logging
 import argparse
 import hashlib
+import os
+import glob
 
 import numpy as np
 
+import gwdatafind
+
 import pycbc
 from pycbc.io import HFile
 
-import gwdatafind
-
 from gwpy.timeseries import TimeSeriesDict
 from gwpy.segments import Segment, DataQualityFlag
 
+
+def find_frames(ifo, frame_dir, frame_type, start, end):
+    """
+    Find the frame files for a given time range.
+    """
+    ifo_frame_type = frame_type.format(ifo=ifo)
+    ifo_pattern = f'{ifo[0]}-{ifo_frame_type}_llhoft-'
+    folder_pattern = ifo_pattern + '{supertime}'
+    path_pattern = os.path.join(
+        frame_dir,
+        ifo_frame_type,
+        folder_pattern,
+    )
+
+    start_round = int(start / 10000)
+    end_round = int(end / 10000)
+    files = []
+    for t in range(start_round, end_round + 1):
+        framedir = path_pattern.format(supertime=t)
+        files += glob.glob(os.path.join(framedir, '*.gwf'))
+
+    return files
+
+
 parser = argparse.ArgumentParser(description=__doc__)
 parser.add_argument("--trigger-file", required=True)
 parser.add_argument("--ifo", required=True)
 parser.add_argument("--gps-start-time", required=True, type=int)
 parser.add_argument("--gps-end-time", required=True, type=int)
 parser.add_argument("--template-bin-file", required=True)
-parser.add_argument("--analysis-frame-type", required=True)
+parser.add_argument("--use-gwdatafind", action='store_true',
+                    help='Use gwdatafind to find frame files. If provided, '
+                         'frame-directory argument will be ignored.')
+parser.add_argument("--frame-directory",
+                    help='Directory containing frame files. Required if '
+                         'not using gwdatafind.')
+parser.add_argument("--frame-type", required=True)
 parser.add_argument("--analysis-flag-name", required=True)
 parser.add_argument("--dq-channel", required=True)
 parser.add_argument("--dq-ok-channel", required=True)
 parser.add_argument("--dq-thresh", required=True, type=float)
+parser.add_argument("--dq-padding", type=float, default=1.0)
+parser.add_argument("--replay-offset", type=int, default=0)
 parser.add_argument("--output", required=True)
 pycbc.add_common_pycbc_options(parser)
 args = parser.parse_args()
 
 pycbc.init_logging(args.verbose)
 
+if not args.use_gwdatafind and args.frame_directory is None:
+    raise ValueError('Must provide frame-directory if not using gwdatafind.')
+
 # Get observing segs
 ar_flag_name = args.analysis_flag_name.format(ifo=args.ifo)
 day_seg = Segment(args.gps_start_time, args.gps_end_time)
 observing_flag = DataQualityFlag.query(ar_flag_name, day_seg)
+
+# shift observing flag to current time
+observing_flag = observing_flag.pad(args.replay_offset, args.replay_offset)
+
 observing_segs = observing_flag.active
 livetime = observing_flag.livetime
 logging.info(f'Found {livetime} seconds of observing time at {args.ifo}.')
 
 # for each segment, check how much time was dq flagged
 flagged_time = 0
-frame_type = args.analysis_frame_type.format(ifo=args.ifo)
 dq_channel = args.dq_channel.format(ifo=args.ifo)
 dq_ok_channel = args.dq_ok_channel.format(ifo=args.ifo)
 for seg in observing_segs:
-    frames = gwdatafind.find_urls(
-        args.ifo[0],
-        frame_type,
-        seg[0],
-        seg[1],
-    )
+    if args.use_gwdatafind:
+        frame_type = args.analysis_frame_type.format(ifo=args.ifo)
+        frames = gwdatafind.find_urls(
+            args.ifo[0],
+            frame_type,
+            seg[0],
+            seg[1],
+        )
+    else:
+        frames = find_frames(
+            args.ifo,
+            args.frame_directory,
+            args.frame_type,
+            seg[0],
+            seg[1],
+        )
+
     tsdict = TimeSeriesDict.read(
         frames,
         channels=[dq_channel, dq_ok_channel],
         start=seg[0],
         end=seg[1],
+        pad=0,
     )
+
     dq_ok_flag = (tsdict[dq_ok_channel] == 1).to_dqflag()
     dq_flag = (tsdict[dq_channel] <= args.dq_thresh).to_dqflag()
-    flagged_time += (dq_flag & dq_ok_flag).livetime
+
+    valid_bad_dq = dq_flag & dq_ok_flag
+
+    # pad flag outwards to match the padding used in the live script
+    valid_bad_dq.protract(args.dq_padding)
+    valid_bad_dq.coalesce()
+
+    flagged_time += (valid_bad_dq & observing_flag).livetime
+
 logging.info(f'Found {flagged_time} seconds of dq flagged time at {args.ifo}.')
 
 bg_livetime = livetime - flagged_time

bin/live/pycbc_live_supervise_collated_trigger_fits

@@ -495,10 +495,12 @@ def daily_dq_trigger_rates(
     fname_format = daily_dq_controls['daily-dq-format']
     ddtr_out_fname = fname_format.format(date=day_str, ifo=ifo)
 
-    today = get_true_date(day_dt, daily_dq_controls)
+    start_time = get_true_date(day_dt, daily_dq_controls)
 
-    gps_start_time = gpstime.utc_to_gps(today).gpsSeconds
-    gps_end_time = gpstime.utc_to_gps(today + timedelta(days=1)).gpsSeconds
+    gps_start_time = gpstime.utc_to_gps(start_time).gpsSeconds
+    gps_end_time = gpstime.utc_to_gps(start_time + timedelta(days=1)).gpsSeconds
+
+    offset = gpstime.utc_to_gps(day_dt).gpsSeconds - gps_start_time
 
     ddtr_out_full = os.path.join(output_dir, ddtr_out_fname)
     daily_dq_args = ['pycbc_live_collated_dq_trigger_rates']
@@ -508,6 +510,7 @@ def daily_dq_trigger_rates(
 
     daily_dq_options['gps-start-time'] = f'{gps_start_time:d}'
     daily_dq_options['gps-end-time'] = f'{gps_end_time:d}'
+    daily_dq_options['replay-offset'] = f'{offset:d}'
     daily_dq_args += sv.dict_to_args(daily_dq_options)
 
     sv.run_and_error(daily_dq_args, controls)