gstlal_inspiral_followups_from_gracedb

#!/usr/bin/env python
#
# Copyright (C) 2012  Kipp Cannon, Chad Hanna, Drew Keppel
#
# 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.

## @file
# A program to request some followup data from a running gstlal_inspiral job based on gracedb submissions notified by lvalert

import sys, os
import logging
from pylal import rate
os.environ["MPLCONFIGDIR"] = "/tmp"
import matplotlib
matplotlib.use('Agg')
import pylab
import numpy
from glue.ligolw import utils
import copy
from gstlal import far
import time
from glue.ligolw import ligolw, table, lsctables, param, array
array.use_in(ligolw.LIGOLWContentHandler)
param.use_in(ligolw.LIGOLWContentHandler)
lsctables.use_in(ligolw.LIGOLWContentHandler)
import urllib
import urlparse
import httplib
import subprocess
from ligo.gracedb.rest import GraceDb
import glob
from ligo.lvalert.utils import get_LVAdata_from_stdin

def get_filename(gracedb_client, graceid, filename, retries = 3, retry_delay = 10.0, ignore_404 = False):
    for i in range(retries):
        logging.info("retrieving \"%s\" for %s" % (filename, graceid))
        response = gracedb_client.files(graceid, filename)
        if response.status == httplib.OK:
            return response
        if response.status == httplib.NOT_FOUND and ignore_404:
            logging.warning("retrieving \"%s\" for %s: (%d) %s.  skipping ..." % (filename, graceid, response.status, response.reason))
            return None
        logging.warning("retrieving \"%s\" for %s: (%d) %s.  pausing ..." % (filename, graceid, response.status, response.reason))
        time.sleep(retry_delay)
    raise Exception("retrieving \"%s\" for %s: (%d) %s" % (filename, graceid, response.status, response.reason))


if len(sys.argv) == 1:
    lvalert_data = get_LVAdata_from_stdin(sys.stdin, as_dict = True)
    logging.info("%(alert_type)s-type alert for event %(uid)s" % lvalert_data)
    logging.info("lvalert data: %s" % repr(lvalert_data))
    filename = os.path.split(urlparse.urlparse(lvalert_data["file"]).path)[-1]
    if filename not in (u"coinc.xml",) and "_CBC_" not in filename:
        logging.info("filename is not 'coinc.xml'.  skipping")
        sys.exit()
    gid = str(lvalert_data["uid"])
else:
    gid = sys.argv[1]

try:
    os.mkdir(gid)
except OSError:
    pass

gracedb = GraceDb()
filename = get_filename(gracedb, gid, "coinc.xml", retries = 3, retry_delay = 10.0, ignore_404 = False)
xmldoc, checksum = utils.load_fileobj(filename, contenthandler = ligolw.LIGOLWContentHandler)
coinc_inspiral_row = lsctables.table.get_table(xmldoc, lsctables.CoincInspiralTable.tableName)[0]
sngl_inspiral_table = lsctables.table.get_table(xmldoc, lsctables.SnglInspiralTable.tableName)
process = lsctables.table.get_table(xmldoc, lsctables.ProcessTable.tableName)
params = lsctables.table.get_table(xmldoc, lsctables.ProcessParamsTable.tableName)
snr_chisq_dict = dict(((ifo,(None,None)) for ifo in ['H1','L1', 'V1']))
for r in sngl_inspiral_table:
    snr_chisq_dict[r.ifo] = (r.snr, r.chisq)
for r in params:
    if r.program == "gstlal_inspiral" and r.param == "--likelihood-file":
        path = r.value
    if r.program == "gstlal_inspiral" and r.param == "--job-tag":
        tag = r.value
for r in process:
    if r.program == "gstlal_inspiral":
        node = r.node

url = open("%s_registry.txt" % tag).readline().strip()

#
# SNR / Time plot
#

#FIXME don't hardcode port number
t,snr = numpy.loadtxt(urllib.urlopen("%ssnr_history.txt" % url), unpack = True)
t -= coinc_inspiral_row.get_end()
pylab.figure()
pylab.semilogy(t, snr)
pylab.ylabel('SNR')
pylab.xlabel('Time from event (s)')
fname = '%s/%s_snrtime.png' % (gid, gid)
pylab.ylim([min(snr), max(snr)])
pylab.xlim([min(t), max(t)])
pylab.grid()
pylab.savefig(fname)
gracedb.writeLog(gid, "SNR vs time", filename = fname, filecontents = open(fname).read(), tagname = "background")

#
# SNR / Chisq plots
#

likelihood_data = far.parse_likelihood_control_doc(utils.load_filename(path, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))[0]

counts = likelihood_data.background_rates
inj = likelihood_data.injection_rates

bgcol = (224/255.,224/255.,224/255.)

likely = copy.deepcopy(inj)
for i, ifo in enumerate(['H1', 'L1', 'V1']):
    snrm, chisqm = snr_chisq_dict[ifo]
    if snrm is None:
        continue
    likely[ifo+"_snr_chi"].array /= counts[ifo+"_snr_chi"].array
    for name, obj in (("background", counts),):
        fig = pylab.figure(figsize=(6,4), facecolor = 'g')
        fig.patch.set_alpha(0.0)
        H1 = obj[ifo+"_snr_chi"].array
        snr = obj[ifo+"_snr_chi"].bins[0].centres()[1:-1]
        chi = obj[ifo+"_snr_chi"].bins[1].centres()[1:-1]
        chi[0] = 0 # not inf
        ax = pylab.subplot(111)
        pylab.pcolormesh(snr, chi, numpy.log10(H1.T +1)[1:-1,1:-1])
        if snrm is not None and chisqm is not None:
            pylab.plot(snrm, chisqm / snrm / snrm, 'ko', mfc = 'None', ms = 14, mew=4)
        if "Log" in str(obj[ifo+"_snr_chi"].bins[0]):
            ax.set_xscale('log')
        if "Log" in str(obj[ifo+"_snr_chi"].bins[1]):
            ax.set_yscale('log')
        pylab.colorbar()
        pylab.xlabel('SNR')
        pylab.ylabel('reduced chi^2 / SNR^2')
        pylab.ylim([0.001, 1])
        pylab.xlim([4, 200])
        pylab.title('%s: %s log base 10 (number + 1)' % (ifo, name))
        pylab.grid(color=(0.1,0.4,0.5), linewidth=2)
        fname = '%s/%s_%s_snrchi.png' % (gid, gid, ifo)
        pylab.savefig(fname)
        gracedb.writeLog(gid, "%s SNR/Chisq" % ifo, filename = fname, filecontents = open(fname).read(), tagname = "background")

sys.exit()

#
# QScans
# FIXME currently broken
#

for r in sngl_inspiral_table:
    cachefile = "%s/%sframes.cache" % (gid, r.ifo)
    subprocess.call(["createframecache.pl", cachefile, "/scratch/llcache/llhoft/%s" % r.ifo])
    conffile = "%s/%sconfig.txt" % (gid, r.ifo)
    f = open(conffile , "w")
    channel = "%s:%s" % (r.ifo, r.channel)
    f.write("""[%s,%s]

{
   channelName:                 '%s'
   frameType:                   '%s_llhoft'
   sampleFrequency:             16384
   searchFrequencyRange:        [10 2048]
   #searchChirpRange:            [+0 +0]
   searchTimeRange:              64
   searchQRange:                [4 64]
   searchMaximumEnergyLoss:     0.2
   searchWindowDuration:        0.25
   whiteNoiseFalseRate:         1e-3
   plotTimeRanges:              10
   plotFrequencyRange:          []
   plotNormalizedEnergyRange:   [0 25.5]
   alwaysPlotFlag:               1
} """ % (channel, channel, channel, r.ifo))
    f.close()

    tstr = str(r.get_end())

    # actually run the scan
    subprocess.call(["dmt_wscan", str(r.get_end()), conffile, cachefile, gid, "1", "1"])

# upload some plots to gracedb
for f in glob.glob("%s/*_spectrogram_autoscaled.png" % gid):
    gracedb.writeLog(gid, "DMT omega scan", filename = f, filecontents = open(f).read(), tagname = "background")
    
sys.exit()