From 7a7fb02fb87ba1a0ada095280c49ed2643461502 Mon Sep 17 00:00:00 2001
From: Tito Dal Canton <tito.dalcanton@ijclab.in2p3.fr>
Date: Sun, 18 Feb 2024 17:15:58 +0100
Subject: [PATCH] Remove the need for template_duration from
 plot_singles_timefreq

---
 bin/plotting/pycbc_plot_singles_timefreq | 35 ++++++++++++++++--------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/bin/plotting/pycbc_plot_singles_timefreq b/bin/plotting/pycbc_plot_singles_timefreq
index e8d12e9575a..6df71f23de8 100644
--- a/bin/plotting/pycbc_plot_singles_timefreq
+++ b/bin/plotting/pycbc_plot_singles_timefreq
@@ -97,13 +97,31 @@ pc = ax.pcolormesh(t + opts.gps_start_time - center_time, freq, Pxx,
                    norm=norm, cmap='afmhot_r', shading='gouraud')
 del freq, Pxx
 
+logging.info('Loading bank')
+bank = pycbc.waveform.bank.TemplateBank(
+    opts.bank_file,
+    approximant=opts.approximant,
+    parameters=[
+        'mass1',
+        'mass2',
+        'spin1z',
+        'spin2z',
+        'approximant',
+        'template_duration'
+    ],
+    load_compressed=False
+)
+tmplts = bank.table
+
 logging.info('Loading trigs')
 trig_f = HFile(opts.trig_file, 'r')
 trigs = trig_f[opts.detector]
 
-def rough_filter(snr, chisq, chisq_dof, end_time, tmp_id, tmp_dur):
-    return np.logical_and(end_time > opts.gps_start_time,
-                          end_time < opts.gps_end_time + tmp_dur)
+def rough_filter(_, _, _, end_time, tmp_id):
+    return np.logical_and(
+        end_time > opts.gps_start_time,
+        end_time < opts.gps_end_time + tmplts.template_duration[tmp_id]
+    )
 
 indices, data_tuple = trig_f.select(
     rough_filter,
@@ -112,10 +130,10 @@ indices, data_tuple = trig_f.select(
     'chisq_dof',
     'end_time',
     'template_id',
-    'template_duration',
     group=opts.detector
 )
-snr, chisq, chisq_dof, end_time, template_ids, template_duration = data_tuple
+snr, chisq, chisq_dof, end_time, template_ids = data_tuple
+del data_tuple
 
 if len(indices) > 0:
     if opts.veto_file:
@@ -151,13 +169,6 @@ if len(indices) > 0:
     except ValueError:
         max_rank = None
 
-    logging.info('Loading bank')
-    bank = pycbc.waveform.bank.TemplateBank(opts.bank_file,
-            approximant=opts.approximant,
-            parameters=['mass1', 'mass2', 'spin1z', 'spin2z', 'approximant'],
-            load_compressed=False)
-    tmplts = bank.table
-
     logging.info('Plotting %d trigs', len(sorted_end_time))
     for tc, rho, tid in zip(sorted_end_time, sorted_rank, sorted_template_ids):
         track_t, track_f = pycbc.pnutils.get_inspiral_tf(