Commit 79d9542a authored by Leo Pound Singer's avatar Leo Pound Singer

Teach bayestar_localize_coincs to work with pycbc

Original: 036d53a5a364e23c5ba41de8d3481bb6ba3015f0
parent 2358ba63
......@@ -58,9 +58,9 @@ parser = command.ArgumentParser(
parents=[command.waveform_parser, command.prior_parser, command.skymap_parser])
parser.add_argument('--keep-going', '-k', default=False, action='store_true',
help='Keep processing events if a sky map fails to converge [default: no]')
parser.add_argument('input', metavar='INPUT.xml[.gz]', default='-', nargs='?',
parser.add_argument('input', metavar='INPUT.xml[.gz]', default='-', nargs='+',
type=argparse.FileType('rb'),
help='Input LIGO-LW XML file [default: stdin]')
help='Input LIGO-LW XML file [default: stdin] or PyCBC HDF5 files. If PyCBC files, must be bank file, coinc file, and trigger files, in that order.')
parser.add_argument('--psd-files', nargs='*',
help='pycbc-style merged HDF5 PSD files')
parser.add_argument('--coinc-event-id', type=int, nargs='*',
......@@ -95,19 +95,31 @@ from lalinference.bayestar.sky_map import ligolw_sky_map, rasterize
import os
import sys
import numpy as np
import h5py
# Read coinc file.
log.info('%s:reading input XML file', opts.input.name)
xmldoc, _ = ligolw_utils.load_fileobj(
opts.input, contenthandler=ligolw_bayestar.LSCTablesAndSeriesContentHandler)
log.info('%s:reading input files', ','.join(file.name for file in opts.input))
if os.path.splitext(opts.input[0].name)[1] in {'.hdf', '.hdf5', '.h5'}:
hdf5files = [h5py.File(file.name, 'r') for file in opts.input]
snr_dict = {}
coinc_and_sngl_inspirals = ligolw_bayestar.coinc_and_sngl_inspirals_for_hdf(*hdf5files)
else:
infile, = opts.input
xmldoc, _ = ligolw_utils.load_fileobj(
infile,
contenthandler=ligolw_bayestar.LSCTablesAndSeriesContentHandler)
snr_dict = ligolw_bayestar.snr_series_by_sngl_inspiral_id_for_xmldoc(xmldoc)
coinc_and_sngl_inspirals = ligolw_bayestar.coinc_and_sngl_inspirals_for_xmldoc(xmldoc)
command.mkpath(opts.output)
if opts.condor_submit:
if opts.coinc_event_id:
raise ValueError('must not set --coinc-event-id with --condor-submit')
coinc_event_ids = [int(coinc.coinc_event_id) for coinc, _ in
ligolw_bayestar.coinc_and_sngl_inspirals_for_xmldoc(xmldoc)]
cmd = ['condor_submit', 'accounting_group=ligo.dev.o3.cbc.pe.bayestar',
'on_exit_remove = (ExitBySignal == False) && (ExitCode == 0)',
'on_exit_hold = (ExitBySignal == True) || (ExitCode != 0)',
......@@ -118,7 +130,7 @@ if opts.condor_submit:
'arguments="-B ' + ' '.join(arg for arg in sys.argv
if arg != '--condor-submit') + ' --coinc-event-id $(CoincEventId)"',
'-append', 'queue CoincEventId in ' + ' '.join(
str(coinc_event_id) for coinc_event_id in coinc_event_ids),
str(coinc_event_id) for coinc_event_id in coinc_and_sngl_inspirals),
'/dev/null']
os.execvp('condor_submit', cmd)
......@@ -187,20 +199,21 @@ else:
reference_psd_filenames_by_process_id[sngl_inspiral.process_id])
for sngl_inspiral in sngl_inspirals)
snr_dict = ligolw_bayestar.snr_series_by_sngl_inspiral_id_for_xmldoc(xmldoc)
count_sky_maps_failed = 0
# Loop over all coinc_event <-> sim_inspiral coincs.
for coinc, sngl_inspirals in ligolw_bayestar.coinc_and_sngl_inspirals_for_xmldoc(xmldoc):
if opts.coinc_event_id:
coinc_and_sngl_inspirals = {
coinc_event_id: coinc_and_sngl_inspirals[coinc_event_id]
for coinc_event_id in opts.coinc_event_id}
if opts.coinc_event_id and int(coinc.coinc_event_id) not in opts.coinc_event_id:
continue
for int_coinc_event_id, sngl_inspirals in coinc_and_sngl_inspirals.items():
coinc_event_id = 'coinc_event:coinc_event_id:{}'.format(int_coinc_event_id)
instruments = {sngl_inspiral.ifo for sngl_inspiral in sngl_inspirals}
# Look up PSDs
log.info('%s:reading PSDs', coinc.coinc_event_id)
log.info('%s:reading PSDs', coinc_event_id)
psds = reference_psds_for_sngls(sngl_inspirals)
# Look up SNR time series
......@@ -211,9 +224,9 @@ for coinc, sngl_inspirals in ligolw_bayestar.coinc_and_sngl_inspirals_for_xmldoc
# Loop over sky localization methods
for method in opts.method:
log.info("%s:method '%s':computing sky map", coinc.coinc_event_id, method)
log.info("%s:method '%s':computing sky map", coinc_event_id, method)
if opts.chain_dump:
chain_dump = '%s.chain.npy' % int(coinc.coinc_event_id)
chain_dump = '%s.chain.npy' % int_coinc_event_id
else:
chain_dump = None
try:
......@@ -223,15 +236,15 @@ for coinc, sngl_inspirals in ligolw_bayestar.coinc_and_sngl_inspirals_for_xmldoc
psds=psds, method=method, nside=opts.nside,
chain_dump=chain_dump, phase_convention=opts.phase_convention,
snr_series=snrs, enable_snr_series=opts.enable_snr_series)
sky_map.meta['objid'] = str(coinc.coinc_event_id)
sky_map.meta['objid'] = coinc_event_id
except (ArithmeticError, ValueError):
log.exception("%s:method '%s':sky localization failed", coinc.coinc_event_id, method)
log.exception("%s:method '%s':sky localization failed", coinc_event_id, method)
count_sky_maps_failed += 1
if not opts.keep_going:
raise
else:
log.info("%s:method '%s':saving sky map", coinc.coinc_event_id, method)
filename = '%s.%s.fits' % (int(coinc.coinc_event_id), method)
log.info("%s:method '%s':saving sky map", coinc_event_id, method)
filename = '%d.%s.fits' % (int_coinc_event_id, method)
fits.write_sky_map(os.path.join(opts.output, filename),
sky_map, nest=True)
......
#
# Copyright (C) 2013-2016 Leo Singer
# 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
......@@ -24,6 +24,7 @@ __author__ = "Leo Singer <leo.singer@ligo.org>"
# Python standard library imports.
import itertools
import operator
import collections
# LIGO-LW XML imports.
from glue.ligolw.utils import process as ligolw_process
......@@ -32,8 +33,13 @@ from glue.ligolw import array as ligolw_array
from glue.ligolw import param as ligolw_param
from glue.ligolw import table as ligolw_table
from glue.ligolw import lsctables
import lal
import lal.series
# Third-party imports.
import numpy as np
from astropy.table import Table
# FIXME: Copied from pylal.ligolw_thinca to avoid dependency.
# Should be moved to lalinspiral.
......@@ -127,37 +133,108 @@ def sim_coinc_and_sngl_inspirals_for_xmldoc(xmldoc):
yield sim_inspiral, coinc, sngl_inspirals
def coinc_and_sngl_inspirals_for_xmldoc(xmldoc, coinc_def=InspiralCoincDef):
class coinc_and_sngl_inspirals_for_hdf(collections.Mapping):
def __init__(self, bank_file, coinc_file, *trigger_files, **kwargs):
self.bank = tuple(bank_file.values())
sample = kwargs.get('sample', 'foreground')
key_prefix = 'detector_'
template_nums, self.ifos = zip(*((key[len(key_prefix):], value)
for key, value in coinc_file.attrs.items()
if key.startswith(key_prefix)))
coinc_group = coinc_file[sample]
self.template_ids = coinc_group['template_id']
self.trigger_ids = [
coinc_group['trigger_id' + template_num]
for template_num in template_nums]
triggers = {}
for f in trigger_files:
(ifo, group), = f.items()
triggers[ifo] = [
group['snr'], group['coa_phase'], group['end_time']]
self.triggers = tuple(triggers[ifo] for ifo in self.ifos)
# Declare types for ersatz CoincEventTable and SnglInspiralTable rows.
self.SnglInspiral = collections.namedtuple('SnglInspiral',
['event_id', 'ifo', 'snr', 'coa_phase', 'end'] + bank_file.keys())
def __getitem__(self, coinc_id):
template_id = self.template_ids[coinc_id]
template = [col[template_id] for col in self.bank]
trigger_ids = [
trigger_ids[coinc_id] for trigger_ids in self.trigger_ids]
return [
self.SnglInspiral(
trigger_id, ifo,
triggers[0][trigger_id], triggers[1][trigger_id],
lal.LIGOTimeGPS(triggers[2][trigger_id]), *template
) for trigger_id, ifo, triggers
in zip(trigger_ids, self.ifos, self.triggers)
]
def __iter__(self):
return iter(range(len(self)))
def __len__(self):
return len(self.template_ids)
class coinc_and_sngl_inspirals_for_xmldoc(collections.Mapping):
"""Retrieve (as a generator) all of the
(sngl_inspiral, sngl_inspiral, ... sngl_inspiral) tuples from coincidences
in a LIGO-LW XML document."""
# Look up necessary tables.
coinc_table = ligolw_table.get_table(xmldoc, lsctables.CoincTable.tableName)
coinc_def_table = ligolw_table.get_table(xmldoc, lsctables.CoincDefTable.tableName)
coinc_map_table = ligolw_table.get_table(xmldoc, lsctables.CoincMapTable.tableName)
sngl_inspiral_table = ligolw_table.get_table(xmldoc, lsctables.SnglInspiralTable.tableName)
# Look up coinc_def id.
if coinc_def:
sngl_sngl_coinc_def_ids = {row.coinc_def_id for row in coinc_def_table
if (row.search, row.search_coinc_type) ==
(coinc_def.search, coinc_def.search_coinc_type)}
# Indices to speed up lookups by ID.
key = operator.attrgetter('coinc_event_id')
coinc_maps_by_coinc_event_id = {coinc_event_id: tuple(coinc_maps)
for coinc_event_id, coinc_maps
in itertools.groupby(sorted(coinc_map_table, key=key), key=key)}
sngl_inspirals_by_event_id = {sngl_inspiral.event_id: sngl_inspiral
for sngl_inspiral in sngl_inspiral_table}
# Loop over all sngl_inspiral <-> sngl_inspiral coincs.
for coinc in coinc_table:
if not coinc_def or coinc.coinc_def_id in sngl_sngl_coinc_def_ids:
coinc_maps = coinc_maps_by_coinc_event_id[coinc.coinc_event_id]
yield coinc, tuple(sngl_inspirals_by_event_id[coinc_map.event_id]
for coinc_map in coinc_maps)
def __init__(self, xmldoc, coinc_def=InspiralCoincDef):
# Look up necessary tables.
coinc_map_table = ligolw_table.get_table(
xmldoc, lsctables.CoincMapTable.tableName)
sngl_inspiral_table = ligolw_table.get_table(
xmldoc, lsctables.SnglInspiralTable.tableName)
# Indices to speed up lookups by ID.
key = operator.attrgetter('coinc_event_id')
self.coinc_maps_by_coinc_event_id = {
int(coinc_event_id):
tuple(int(coinc_map.event_id) for coinc_map in coinc_maps)
for coinc_event_id, coinc_maps
in itertools.groupby(sorted(coinc_map_table, key=key), key=key)}
self.sngl_inspirals_by_event_id = {
int(row.event_id): row for row in sngl_inspiral_table}
# Look up coinc_def_ids.
if coinc_def is None:
self.coinc_event_ids = sorted({
int(row.coinc_event_id) for row in coinc_map_table})
else:
coinc_table = ligolw_table.get_table(
xmldoc, lsctables.CoincTable.tableName)
coinc_def_table = ligolw_table.get_table(
xmldoc, lsctables.CoincDefTable.tableName)
coinc_def_ids = {
row.coinc_def_id for row in coinc_def_table
if (row.search, row.search_coinc_type) ==
(coinc_def.search, coinc_def.search_coinc_type)}
self.coinc_event_ids = tuple(
int(row.coinc_event_id) for row in coinc_table
if row.coinc_def_id in coinc_def_ids)
def __getitem__(self, coinc_event_id):
coinc_maps = self.coinc_maps_by_coinc_event_id[coinc_event_id]
return tuple(
self.sngl_inspirals_by_event_id[event_id]
for event_id in coinc_maps)
def __iter__(self):
return iter(self.coinc_event_ids)
def __len__(self):
return len(self.coinc_event_ids)
def psd_filenames_by_process_id_for_xmldoc(xmldoc):
......
......@@ -417,8 +417,8 @@ def gracedb_sky_map(
phase_convention = 'antifindchirp'
# Locate the sngl_inspiral rows that we need.
(_, sngl_inspirals), = ligolw.coinc_and_sngl_inspirals_for_xmldoc(
xmldoc, coinc_def=None)
sngl_inspirals, = ligolw.coinc_and_sngl_inspirals_for_xmldoc(
xmldoc, coinc_def=None).values()
sngl_inspirals = list(sngl_inspirals)
# Try to load complex SNR time series.
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment