gstlal_inspiral_compute_dtdphideff_cov_matrix

#!/usr/bin/python
#
# Copyright (C) 2019 Chad Hanna, Amit Reza, Leo Tsukada
#
# 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 2 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.
import argparse
import itertools
import numpy, scipy.interpolate
from lal import series
from gstlal.stats.inspiral_extrinsics import TimePhaseSNR
from gstlal import templates
from ligo.lw import utils as ligolw_utils
import sys
import h5py

parser = argparse.ArgumentParser(description = 'generate a dt dphi covariance matrix')
parser.add_argument('--psd-xml', help = 'XML containing HLV psd')
parser.add_argument('--H-snr', type = float, help = 'H characteristic SNR')
parser.add_argument('--L-snr', type = float, help = 'L characteristic SNR')
parser.add_argument('--V-snr', type = float, help = 'V characteristic SNR')
parser.add_argument('--K-snr', type = float, help = 'K characteristic SNR')
parser.add_argument('--m1', type = float, default = 1.4, help = 'primary component mass')
parser.add_argument('--m2', type = float, default = 1.4, help = 'secondary component mass')
parser.add_argument('--s1', type = float, default = 0., help = 'primary (z) spin')
parser.add_argument('--s2', type = float, default = 0., help = 'secondary (z) spin')
parser.add_argument('--flow', type = float, default = 10., help = 'Low frequency cut-off. Default 10 Hz')
parser.add_argument('--fhigh', type = float, default = 1024., help = 'High frequency cut-off. Default 1024 Hz')
parser.add_argument("--output", help = "set the output h5 file, e.g., covmat.h5")
args = parser.parse_args()

if args.flow >= args.fhigh:
	raise ValueError("flow cannot be greater than fhigh")

refpsd = args.psd_xml
rho = {"L1": args.L_snr, "H1": args.H_snr, "V1": args.V_snr, "K1": args.K_snr}


psd = series.read_psd_xmldoc(ligolw_utils.load_filename(refpsd, verbose = True, contenthandler = series.PSDContentHandler))
psddict = {}
for ifo in rho:
	ifo_psd = ifo
	# FIXME use Virgo PSD for KAGRA temporarily. This has to be fixed once KAGRA PSD is measured.
	if rho == "K1":
		ifo = "V1"
	f = numpy.arange(len(psd[ifo].data.data)) * psd[ifo].deltaF
	psddict[ifo_psd] = scipy.interpolate.interp1d(f, psd[ifo].data.data)

def moment(fmin, fmax, n, m1, m2, s1, s2, ifo, psddict = psddict, delta_f = 0.25):
	farr = numpy.linspace(0, fmax, fmax / delta_f + delta_f)
	h = templates.hplus_of_f(m1, m2, s1, s2, fmin, fmax, delta_f)
	return templates.moment(farr, h, n, psddict[ifo])

sigsqtt = {}
sigsqpp = {}
sigsqtp = {}
sigsqdd = {}

for ifo in rho:
	sigsqf = moment(args.flow, args.fhigh, 2, args.m1, args.m2, args.s1, args.s2, ifo) - moment(args.flow, args.fhigh, 1, args.m1, args.m2, args.s1, args.s2, ifo)**2
	sigsqtt[ifo] = (1. / (2 * 3.14 * rho[ifo] * sigsqf**.5)**2)
	sigsqpp[ifo] = moment(args.flow, args.fhigh, 2, args.m1, args.m2, args.s1, args.s2, ifo) / (rho[ifo]**2 * sigsqf)
	sigsqtp[ifo] = moment(args.flow, args.fhigh, 1, args.m1, args.m2, args.s1, args.s2, ifo) / (2 * 3.14 * rho[ifo]**2 * sigsqf)
	sigsqdd[ifo] = 1. / rho[ifo]**2

transtt = {}
transpp = {}
transtp = {}
transpt = {}
transdd = {}
# FIXME do an actual calculation
# transdd = {frozenset(("H1", "L1")): 1. / numpy.sqrt(sigsqdd["H1"] + sigsqdd["L1"])  , frozenset(("H1", "V1")): 1. / numpy.sqrt(sigsqdd["H1"] + sigsqdd["V1"]), frozenset(("L1", "V1")): 1. / numpy.sqrt(sigsqdd["L1"] + sigsqdd["V1"])}
# transdd = {frozenset(instrument_pair) : 1. / numpy.sqrt(sigsqdd[instrument_pair[0]] + sigsqdd[instrument_pair[1]]) for instrument_pair in TimePhaseSNR.instument_pairs(("H1", "L1", "V1", "K1"))}

for pair in TimePhaseSNR.instument_pairs(("H1", "L1", "V1", "K1")):
	a,b = pair
	m11 = sigsqtt[a] + sigsqtt[b]
	m22 = sigsqpp[a] + sigsqpp[b]
	m12 = m21 = sigsqtp[a] + sigsqtp[b]
	mat = numpy.array([[m11, m12], [m21, m22]])
	matinv = numpy.linalg.inv(mat)
	cholesky_transpose = numpy.linalg.cholesky(matinv).T
	transtt[frozenset(pair)] = cholesky_transpose[0,0]
	transtp[frozenset(pair)] = cholesky_transpose[0,1]
	transpt[frozenset(pair)] = cholesky_transpose[1,0]
	transpp[frozenset(pair)] = cholesky_transpose[1,1]
	transdd[frozenset(pair)] = 1. / numpy.sqrt(sigsqdd[a] + sigsqdd[b])

f = h5py.File(args.output, "w")
h5_transtt = f.create_group("transtt")
h5_transtp = f.create_group("transtp")
h5_transpt = f.create_group("transpt")
h5_transpp = f.create_group("transpp")
h5_transdd = f.create_group("transdd")
for group, mat in zip((h5_transtt, h5_transtp, h5_transpt, h5_transpp, h5_transdd), (transtt, transtp, transpt, transpp, transdd)):
	for k,v in mat.items():
		group.create_dataset(",".join(sorted(k)), data = float(v))

combos = TimePhaseSNR.instrument_combos(("H1","L1","V1", "K1")) + (("H1",),("L1",),("V1",), ("K1",))
norm = dict((frozenset(k), 0.) for k in combos)
norm = {}
h5_norm = f.create_group("norm")
for k in combos:
	h5_norm.create_dataset(",".join(sorted(k)), data = 1.0)
	norm[k] = 1.0
f.close()

print "transtt =", transtt
print "transtp =", transtp
print "transpt =", transpt
print "transpp =", transpp
print "transdd =", transdd
print "norm = ", norm