#!/usr/bin/env python
# Copyright (C) 2015 Alexander Harvey Nitz
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" Calculate the SNR and CHISQ timeseries for either a chosen template, or
a specific Nth loudest coincident event.
"""
import sys, logging, argparse, numpy, pycbc, h5py
from pycbc import vetoes, psd, waveform, strain, scheme, fft, filter
from pycbc.io import WaveformArray
from pycbc import events
from pycbc.filter import resample_to_delta_t
from pycbc.types import zeros, complex64
from pycbc.types import complex_same_precision_as
import pycbc.waveform.utils
import pycbc.version
def subtract_template(stilde, template, snr, trigger_time, flow):
idx = int((trigger_time - snr.start_time) / snr.delta_t)
if idx >= 0 and idx < len(snr):
sig = filter.sigma(template, psd=stilde.psd, low_frequency_cutoff=flow)
inverse = template * snr[idx] / sig
dt = trigger_time - snr.start_time
stilde -= pycbc.waveform.utils.apply_fseries_time_shift(inverse, dt)
return stilde
def select_segments(fname, anal_name, data_name, ifo, time, pad_data):
anal_segs = events.select_segments_by_definer(fname, anal_name, ifo)
data_segs = events.select_segments_by_definer(fname, data_name, ifo)
# Anal segs should be disjoint, so first find the seg containing time
s = numpy.array([t[0] for t in anal_segs])
e = numpy.array([t[1] for t in anal_segs])
#ensure sorted
sorting = s.argsort()
s = s[sorting]
e = e[sorting]
idx = numpy.searchsorted(s, time) - 1
anal_time = (s[idx], e[idx])
# Now need to find the corresponding data_seg. This could be complicated
# as in edge cases the anal_time tuple could be completely contained within
# *two* data blocks (think analysis chunk slightly longer than minimum).
# We need to choose the *right* block to reproduce what the search did.
s2 = numpy.array([t[0] for t in data_segs])
e2 = numpy.array([t[1] for t in data_segs])
lgc = (s2 < time) & (e2 > time)
s2 = s2[lgc]
e2 = e2[lgc]
if len(s2) == 0:
err_msg = "Cannot find a data segment within %s" %(str(time))
raise ValueError(err_msg)
if len(s2) == 1:
data_time = (s2[0], e2[0])
if len(s2) > 1:
# The tricksy case. The corresponding data segment should have the
# largest overlap with anal_time
overlap = None
for start, end in zip(s2, e2):
if start + pad_data > anal_time[0] or end - pad_data < anal_time[1]:
# The analysis time must lie within the data time, otherwise
# this clearly is not the right data segment!
continue
curr_nonoverlap = anal_time[0] - start
curr_nonoverlap += end - anal_time[1]
if (overlap is None) or (curr_nonoverlap < overlap):
overlap = curr_nonoverlap
data_time = (start, end)
return anal_time, data_time
parser = argparse.ArgumentParser(usage='',
description="Single template gravitational-wave followup")
parser.add_argument('--version', action=pycbc.version.Version)
parser.add_argument('--output-file', required=True)
parser.add_argument('--subtract-template', action='store_true')
parser.add_argument("-V", "--verbose", action="store_true",
help="print extra debugging information", default=False )
parser.add_argument("--low-frequency-cutoff", type=float,
help="The low frequency cutoff to use for filtering (Hz)")
parser.add_argument("--chisq-bins", default="0", type=str, help=
"Number of frequency bins to use for power chisq.")
parser.add_argument("--minimum-chisq-bins", default=0, type=int, help=
"If the chisq bin formula fails, default to this number.")
parser.add_argument("--trigger-time", type=float, default=0,
help="The central time of the trigger to use.")
parser.add_argument("--use-params-of-closest-injection", action="store_true",
default=False,
help="If given, use the injection with end_time closest to "
"--trigger-time as the waveform for filtering. If "
"using this do not supply mass and spin options. "
"Using this requires trigger-time and injection-file "
"to be given.")
# add approximant arg
waveform.bank.add_approximant_arg(parser,
help="The name of the approximant to use for filtering. "
"Do not use if using --use-params-of-closest-injection.")
parser.add_argument("--mass1", type=float,
help="The mass of the first component object. "
"Do not use if using --use-params-of-closest-injection.")
parser.add_argument("--mass2", type=float,
help="The mass of the second component object. "
"Do not use if using --use-params-of-closest-injection.")
parser.add_argument("--spin1z", type=float, default=0,
help="The aligned spin of the first component object. "
"Do not use if using --use-params-of-closest-injection.")
parser.add_argument("--spin2z", type=float, default=0,
help="The aligned pin of the second component object. "
"Do not use if using --use-params-of-closest-injection.")
parser.add_argument("--template-start-frequency", type=float, default=None,
help="If given, use this as a start frequency for "
"generating the template. If not given the "
"--low-frequency-cutoff is used.")
# Optional arguments for precessing templates
parser.add_argument("--spin1x", type=float, default=0,
help="The non-aligned spin of the first component object. "
"Default = 0.")
parser.add_argument("--spin2x", type=float, default=0,
help="The non-aligned spin of the second component object. "
"Default = 0.")
parser.add_argument("--spin1y", type=float, default=0,
help="The non-aligned spin of the first component object. "
"Default = 0.")
parser.add_argument("--spin2y", type=float, default=0,
help="The non-aligned spin of the second component object. "
"Default = 0.")
parser.add_argument("--inclination", type=float, default=0,
help="The inclination of the source w.r.t the observer. "
"Default = 0.")
parser.add_argument("--u-val", type=float, default=None,
help="The ratio between hplus and hcross to use in the "
"template, according to h(t) = hplus * u_val + hcross. "
"If not given only hplus is used.")
parser.add_argument("--window", type=float,
help="Time to save around the trigger time (if given)")
parser.add_argument("--order", type=int,
help="The integer half-PN order at which to generate"
" the approximant. Default is -1 which indicates to use"
" approximant defined default.", default=-1,
choices = numpy.arange(-1, 9, 1))
parser.add_argument("--taper-template", choices=["start","end","startend"],
help="For time-domain approximants, taper the start and/or"
" end of the waveform before FFTing.")
# These options can be used to identify start/end times
parser.add_argument("--inspiral-segments",
help="XML file containing the inspiral analysis segments. "
"Only used with the --statmap-file option")
parser.add_argument("--data-read-name",
help="name of the segmentlist containing the data read in by each job "
"from the inspiral segment file")
parser.add_argument("--data-analyzed-name",
help="name of the segmentlist containing the data analysed by each job "
"from the inspiral segment file")
# Add options groups
psd.insert_psd_option_group(parser)
strain.insert_strain_option_group(parser)
strain.StrainSegments.insert_segment_option_group(parser)
scheme.insert_processing_option_group(parser)
fft.insert_fft_option_group(parser)
opt = parser.parse_args()
f = h5py.File(opt.output_file, 'w')
ifo = opt.channel_name[0:2]
# If we are choosing start/end times from XML file ############################
if opt.inspiral_segments:
# Important for trig-start/end and data-start/end to match inspiral jobs.
# Zero-padding also will not zero-pad unless explicitly told to in the
# trig start/end times
anal_seg, data_seg = select_segments(opt.inspiral_segments,
opt.data_analyzed_name, opt.data_read_name,
ifo, opt.trigger_time, opt.pad_data)
opt.trig_start_time = anal_seg[0]
opt.trig_end_time = anal_seg[1]
opt.gps_start_time = data_seg[0] + opt.pad_data
opt.gps_end_time = data_seg[1] - opt.pad_data
f.attrs['event_time'] = opt.trigger_time
###############################################################################
# Check that the values returned for the options make sense
psd.verify_psd_options(opt, parser)
strain.verify_strain_options(opt, parser)
strain.StrainSegments.verify_segment_options(opt, parser)
scheme.verify_processing_options(opt, parser)
fft.verify_fft_options(opt,parser)
pycbc.init_logging(opt.verbose)
ctx = scheme.from_cli(opt)
gwstrain = strain.from_cli(opt, pycbc.DYN_RANGE_FAC)
strain_segments = strain.StrainSegments.from_cli(opt, gwstrain)
if not opt.use_params_of_closest_injection:
row = WaveformArray.from_kwargs(
mass1=opt.mass1,
mass2=opt.mass2,
spin1x=opt.spin1x,
spin1y=opt.spin1y,
spin1z=opt.spin1z,
spin2x=opt.spin2x,
spin2y=opt.spin2y,
spin2z=opt.spin2z)
with ctx:
fft.from_cli(opt)
flow = opt.low_frequency_cutoff
# don't want a template f_lower that is None or 0.0
template_flow = opt.template_start_frequency or flow
flen = strain_segments.freq_len
delta_f = strain_segments.delta_f
logging.info("Making frequency-domain data segments")
segments = strain_segments.fourier_segments()
logging.info("Calculating the PSDs")
psd.associate_psds_to_segments(opt, segments, gwstrain, flen, delta_f,
flow, dyn_range_factor=pycbc.DYN_RANGE_FAC,
precision='single')
logging.info("Making template: %s" % opt.approximant)
if opt.use_params_of_closest_injection:
if not hasattr(gwstrain, 'injections') or not opt.trigger_time:
err_msg = "To use --use-params-of-closest-injection you must "
err_msg += "be making injections with an injection file and using "
err_msg += "the --trigger-time option."
raise ValueError(err_msg)
logging.info("Making template directly from injection.")
inj_set = gwstrain.injections
end_times = numpy.array(inj_set.end_times())
inj_idx = abs(end_times - opt.trigger_time).argmin()
inj = inj_set.table[inj_idx]
# Use injection values for things like choosing number of chisq bins
row = WaveformArray.from_kwargs(
mass1=inj.mass1,
mass2=inj.mass2,
spin1x=inj.spin1x,
spin1y=inj.spin1y,
spin1z=inj.spin1z,
spin2x=inj.spin2x,
spin2y=inj.spin2y,
spin2z=inj.spin2z)
# FIXME: Don't like hardcoded 16384 here
# NOTE: f_lower is set in strain.py as inj.f_lower, but this is a
# little unclear to see and caused me some problems! Stating
# this clearly here so no-one makes the same mistake as me.
td_template = inj_set.make_strain_from_inj_object(inj, 1./16384.,
opt.channel_name[0:2], f_lower=inj.f_lower,
distance_scale=opt.injection_scale_factor)
td_template = resample_to_delta_t(td_template, gwstrain.delta_t,
method='ldas')
td_template = td_template * pycbc.DYN_RANGE_FAC
td_template._epoch -= inj.get_end(site=opt.channel_name[0])
approximant = inj.waveform
template = waveform.td_waveform_to_fd_waveform(td_template, length=flen)
template = template.astype(complex_same_precision_as(segments[0]))
else:
approximant = waveform.bank.parse_approximant_arg(opt.approximant, row)[0]
logging.info("Making template: %s" % opt.approximant)
if opt.u_val is None:
template = waveform.get_waveform_filter(
zeros(flen, dtype=complex64),
approximant=approximant,
template=row[0],
inclination=opt.inclination,
taper=opt.taper_template,
f_lower=template_flow, delta_f=delta_f,
delta_t=gwstrain.delta_t,
distance = 1.0/pycbc.DYN_RANGE_FAC)
else:
tp, tc = waveform.get_two_pol_waveform_filter(
zeros(flen, dtype=complex64),
zeros(flen, dtype=complex64), row[0],
approximant=approximant,
inclination=opt.inclination,
taper=opt.taper_template,
f_lower=template_flow, delta_f=delta_f,
delta_t=gwstrain.delta_t)
template = tc.multiply_and_add(tp, opt.u_val)
# FIXME: should probably switch to something like what is done for parsing
# the approximant for chisq bins at some point
class t(object):
pass
parse_row = t()
parse_row.params = t()
for param in row.fieldnames:
setattr(parse_row.params, param, row[param][0])
chisq_bins_float = vetoes.SingleDetPowerChisq.parse_option(parse_row,
opt.chisq_bins)
if numpy.isnan(chisq_bins_float) or (int(chisq_bins_float) <
opt.minimum_chisq_bins):
if opt.minimum_chisq_bins:
chisq_bins = opt.minimum_chisq_bins
logging.warning( "Number of chisq bins is less than minimum or is NaN."
+ (" Using %d bins." % opt.minimum_chisq_bins) )
else:
raise ValueError(
"Chisq bins is NaN or negative and no minimum is set.")
else:
chisq_bins = int(chisq_bins_float)
f['template'] = template.numpy()
snrs, chisqs = [], []
raw_bins = [[] for b in range(chisq_bins)]
if opt.trig_start_time:
start_time = opt.trig_start_time
end_time = opt.trig_end_time
else:
start_time = opt.gps_start_time + opt.segment_start_pad
end_time = opt.gps_end_time - opt.segment_end_pad
if opt.window:
start_time_wind = opt.trigger_time - opt.window
if start_time_wind > start_time:
start_time = start_time_wind
end_time_wind = opt.trigger_time + opt.window
if end_time_wind < end_time:
end_time = end_time_wind
for s_num, stilde in enumerate(segments):
start = stilde.epoch + stilde.analyze.start / float(opt.sample_rate)
end = stilde.epoch + stilde.analyze.stop / float(opt.sample_rate)
if end < start_time:
continue
if start > end_time:
break
logging.info("Filtering segment %s" % s_num)
snr, corr, norm = filter.matched_filter_core(template, stilde,
psd=stilde.psd,
low_frequency_cutoff=flow)
snr *= norm
if opt.subtract_template:
stilde = subtract_template(stilde, template,
snr, opt.trigger_time, flow)
snr, corr, norm = filter.matched_filter_core(template, stilde,
psd=stilde.psd,
low_frequency_cutoff=flow)
logging.info("calculating chisq")
chisq, raw_bin = vetoes.power_chisq(template, stilde, chisq_bins, stilde.psd,
low_frequency_cutoff=flow,
return_bins=True)
chisq /= chisq_bins * 2 - 2
snrs.append(snr[stilde.analyze])
chisqs.append(chisq[stilde.analyze])
for i in range(chisq_bins):
raw_bins[i].append(raw_bin[i][stilde.analyze].numpy())
sidx = int((start_time - snrs[0].start_time) / snr.delta_t)
eidx = int((end_time - start_time) / snr.delta_t) + sidx
if sidx < 0:
err_msg = "Ian has probably broken single_template again. Please email "
err_msg += "with the command line that is raising this error and "
err_msg += "shout at him. Beer may speed up the fixing process."
raise ValueError(err_msg)
for i in range(chisq_bins):
key = 'chisq_bins/%s' % i
f[key] = numpy.concatenate(raw_bins[i])[sidx:eidx]
f[key].attrs['start_time'] = start_time
f[key].attrs['delta_t'] = snr.delta_t
f['chisq_boundaries'] = numpy.array(vetoes.power_chisq_bins(template, chisq_bins, stilde.psd,
low_frequency_cutoff=flow)) * template.delta_f
f['snr'] = numpy.concatenate([snr.numpy() for snr in snrs])[sidx:eidx]
f['snr'].attrs['start_time'] = start_time
f['snr'].attrs['delta_t'] = snr.delta_t
f['chisq'] = numpy.concatenate([chisq.numpy() for chisq in chisqs])[sidx:eidx]
f['chisq'].attrs['start_time'] = start_time
f['chisq'].attrs['delta_t'] = snr.delta_t
f.attrs['approximant'] = approximant
f.attrs['ifo'] = ifo
f.attrs['command_line'] = ' '.join(sys.argv)
logging.info("Finished")