bayestar_localize_lvalert.py 6.07 KB
Newer Older
Adam Mercer's avatar
Adam Mercer committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178
#
# Copyright (C) 2013-2017  Leo Singer
#
# 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.
#
r"""
Listen for new events from LVAlert and perform sky localization.

`bayestar_localize_lvalert` supports two modes of operation. You can
explicitly specify the GraceDb ID on the command line, as in:

    $ bayestar_localize_lvalert T90713

Or, `bayetar_localize_lvalert` can GraceDB IDs from stdin (e.g., from the
terminal, or redirected from a fifo):

    $ mkfifo /var/run/bayestar
    $ tail -F /var/run/bayestar | bayestar_localize_lvalert &
    $ echo T90713 > /var/run/bayestar
"""


#
# Command line interface
#

from lalinference.bayestar import command

methods = '''
    toa_phoa_snr
    toa_phoa_snr_mcmc
    '''.split()
default_method = 'toa_phoa_snr'
command.skymap_parser.add_argument(
    '--method', choices=methods, default=default_method,
    help='Sky localization method [default: %(default)s]')
parser = command.ArgumentParser(
    parents=[
        command.waveform_parser, command.prior_parser, command.skymap_parser])
parser.add_argument(
    '-d', '--disable-detector', metavar='X1', type=str, nargs='+',
    help='disable certain detectors [default: enable all]')
parser.add_argument(
    '-N', '--dry-run', default=False, action='store_true',
    help='Dry run; do not update GraceDB entry [default: %(default)s]')
parser.add_argument(
    '--no-tag', default=False, action='store_true',
    help='Do not set lvem tag for GraceDB entry [default: %(default)s]')
parser.add_argument(
    '-o', '--output', metavar='FILE.fits[.gz]', default='bayestar.fits.gz',
    help='Name for uploaded file [default: %(default)s]')
parser.add_argument(
    'graceid', metavar='G123456', nargs='*',
    help='Run on these GraceDB IDs. If no GraceDB IDs are listed on the '
    'command line, then read newline-separated GraceDB IDs from stdin.')
opts = parser.parse_args()


#
# Late imports
#

import logging
import os
import sys
import six
from lalinference.bayestar.sky_map import localize, rasterize
from lalinference.io import fits
from lalinference.io import events
import ligo.gracedb.logging
import ligo.gracedb.rest

# Squelch annoying and uniformative LAL log messages.
import lal
lal.ClobberDebugLevel(lal.LALNDEBUG)

logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s')
log = logging.getLogger('BAYESTAR')

# If no GraceDB IDs were specified on the command line, then read them
# from stdin line-by-line.
graceids = opts.graceid if opts.graceid else command.iterlines(sys.stdin)

# Fire up a GraceDb client
# FIXME: Mimic the behavior of the GraceDb command line client, where the
# environment variable GRACEDB_SERVICE_URL overrides the default service URL.
# It would be nice to get this behavior into the gracedb package itself.
gracedb = ligo.gracedb.rest.GraceDb(
    os.environ.get(
        'GRACEDB_SERVICE_URL', ligo.gracedb.rest.DEFAULT_SERVICE_URL))

if opts.chain_dump:
    chain_dump = opts.output.replace('.fits.gz', '').replace('.fits', '') + \
        '.chain.npy'
else:
    chain_dump = None

tags = ("sky_loc",)
if not opts.no_tag:
    tags += ("lvem",)

event_source = events.gracedb.open(graceids, gracedb)

if opts.disable_detector:
    event_source = events.detector_disabled.open(
        event_source, opts.disable_detector)

for graceid in six.iterkeys(event_source):

    try:
        event = event_source[graceid]
    except:
        log.exception('failed to read event %s from GraceDB', graceid)
        continue

    # Send log messages to GraceDb too
    if not opts.dry_run:
        handler = ligo.gracedb.logging.GraceDbLogHandler(gracedb, graceid)
        handler.setLevel(logging.INFO)
        logging.root.addHandler(handler)

    # A little bit of Cylon humor
    log.info('by your command...')

    try:
        # perform sky localization
        log.info("starting sky localization")
        sky_map = rasterize(localize(
            event, opts.waveform, opts.f_low, opts.min_distance,
            opts.max_distance, opts.prior_distance_power, opts.cosmology,
            method=opts.method, nside=opts.nside,
            chain_dump=chain_dump, enable_snr_series=opts.enable_snr_series,
            f_high_truncate=opts.f_high_truncate))
        sky_map.meta['objid'] = str(graceid)
        sky_map.meta['url'] = 'https://gracedb.ligo.org/events/{0}'.format(
            graceid)
        log.info("sky localization complete")

        # upload FITS file
        with command.TemporaryDirectory() as fitsdir:
            fitspath = os.path.join(fitsdir, opts.output)
            fits.write_sky_map(fitspath, sky_map, nest=True)
            log.debug('wrote FITS file: %s', opts.output)
            if opts.dry_run:
                command.rename(fitspath, os.path.join('.', opts.output))
            else:
                gracedb.writeLog(
                    graceid, "BAYESTAR rapid sky localization ready",
                    filename=fitspath, tagname=tags)
            log.debug('uploaded FITS file')
    except KeyboardInterrupt:
        # Produce log message and then exit if we receive SIGINT (ctrl-C).
        log.exception("sky localization failed")
        raise
    except:
        # Produce log message for any otherwise uncaught exception.
        # Unless we are in dry-run mode, keep going.
        log.exception("sky localization failed")
        if opts.dry_run:
            # Then re-raise the exception if we are in dry-run mode
            raise

    if not opts.dry_run:
        # Remove old log handler
        logging.root.removeHandler(handler)
        del handler