Add tool to combine sky maps from different observatories

CHANGES.rst

@@ -5,7 +5,8 @@ Changelog
 0.0.8 (unreleased)
 ==================
 
-- No changes yet.
+- Add ``ligo-skymap-combine``, a tool to combine sky localizations from
+  different observations into a joint skymap.
 
 0.0.7 (2018-04-27)
 ==================

docs/index.rst

@@ -92,6 +92,7 @@ Postprocessing
 .. toctree::
    :maxdepth: 1
 
+   ligo/skymap/tool/ligo_skymap_combine
    ligo/skymap/tool/ligo_skymap_from_samples
    ligo/skymap/tool/ligo_skymap_plot_stats
    ligo/skymap/tool/ligo_skymap_stats

docs/ligo/skymap/tool/ligo_skymap_combine.rst

+Joint Sky Localization (`ligo-skymap-combine`)
+==============================================
+
+.. argparse::
+    :module: ligo.skymap.tool.ligo_skymap_combine
+    :func: parser

ligo/skymap/distance.py

@@ -207,6 +207,8 @@ Check that the moments scale as expected when we vary sigma.
 
 Check some more arbitrary values using numerical quadrature:
 
+>>> import scipy.integrate
+>>> from ligo.skymap.core import conditional_pdf
 >>> sigma = 1.0
 >>> for mu in np.linspace(-10, 10):
 ...     mean, std, norm = parameters_to_moments(mu, sigma)

ligo/skymap/tool/ligo_skymap_combine.py

+#
+# Copyright (C) 2018  Tito Dal Canton, Eric Burns, 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 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, see <http://www.gnu.org/licenses/>.
+#
+
+"""Combine different sky localizations of a common event observed by different
+instruments in order to form a more constrained localization.
+
+If one of the input maps contains distance information (for instance from
+BAYESTAR or LALInference) then the marginal distance posterior in the output
+map is updated according to the restriction in sky location imposed by the other
+input map(s). Only one input map can currently have distance information."""
+
+from argparse import FileType
+
+from . import ArgumentParser
+
+
+def parser():
+    parser = ArgumentParser(description=__doc__)
+    parser.add_argument('input', metavar='INPUT.fits[.gz]', type=FileType('rb'),
+                        nargs='+', help='Input sky localizations')
+    # FIXME the output option has type str because astropy.io.fits.writeto()
+    # only honors the .gz extension when given a file name string (as of 3.0.1)
+    parser.add_argument('output', metavar='OUTPUT.fits[.gz]', type=str,
+                        help='Output combined sky localization')
+    parser.add_argument('--origin', type=str,
+                        help='Optional tag describing the organization'
+                             ' responsible for the combined output')
+    return parser
+
+def main(args=None):
+    args = parser().parse_args(args)
+
+    import numpy as np
+    import healpy as hp
+    from astropy.io import fits
+    from astropy.time import Time
+
+    from ..distance import parameters_to_marginal_moments
+    from ..io import read_sky_map, write_sky_map
+
+    input_skymaps = []
+    dist_mu = dist_sigma = dist_norm = None
+    for input_file in args.input:
+        with fits.open(input_file) as hdus:
+            header = hdus[0].header.copy()
+            header.extend(hdus[1].header)
+            nonstd_name = hdus[1].columns.names[0] != 'PROB'
+            if nonstd_name:
+                # FIXME special case for Fermi/GBM localizations which use a
+                # different name for the first field of HDU 1
+                # (PROBABILITY vs PROB)
+                has_distance = False
+                data = (hp.read_map(hdus, nest=True),)
+                meta = {'instruments': set([header['INSTRUME']])}
+            else:
+                has_distance = 'DISTMU' in hdus[1].columns.names
+                data, meta = read_sky_map(hdus, nest=True,
+                                          distances=has_distance)
+
+        nside = hp.npix2nside(len(data[0]))
+        input_skymaps.append((nside, data[0], meta, header))
+
+        if has_distance:
+            if dist_mu is not None:
+                raise RuntimeError('only one input localization can have'
+                                   ' distance information')
+            dist_mu = data[1]
+            dist_sigma = data[2]
+            dist_norm = data[3]
+
+    max_nside = max(x[0] for x in input_skymaps)
+
+    # upsample sky posteriors to maximum resolution and combine them
+    combined_prob = None
+    for nside, prob, _, _ in input_skymaps:
+        if nside < max_nside:
+            prob = hp.ud_grade(prob, max_nside, order_in='NESTED',
+                               order_out='NESTED')
+        if combined_prob is None:
+            combined_prob = np.ones_like(prob)
+        combined_prob *= prob
+
+    # normalize joint posterior
+    norm = combined_prob.sum()
+    if norm == 0:
+        raise RuntimeError('input sky localizations are disjoint')
+    combined_prob /= norm
+
+    out_kwargs = {'gps_creation_time': Time.now().gps,
+                  'nest': True}
+    if args.origin is not None:
+        out_kwargs['origin'] = args.origin
+
+    # average the various input event times
+    input_gps = [x[2]['gps_time'] for x in input_skymaps if 'gps_time' in x[2]]
+    if input_gps:
+        out_kwargs['gps_time'] = np.mean(input_gps)
+
+    # combine instrument tags
+    out_instruments = set()
+    for x in input_skymaps:
+        if 'instruments' in x[2]:
+            out_instruments.update(x[2]['instruments'])
+    out_kwargs['instruments'] = ','.join(out_instruments)
+
+    # update marginal distance posterior, if available
+    if dist_mu is not None:
+        if hp.npix2nside(len(dist_mu)) < max_nside:
+            dist_mu = hp.ud_grade(dist_mu, max_nside, order_in='NESTED',
+                                  order_out='NESTED')
+            dist_sigma = hp.ud_grade(dist_sigma, max_nside, order_in='NESTED',
+                                     order_out='NESTED')
+            dist_norm = hp.ud_grade(dist_norm, max_nside, order_in='NESTED',
+                                    order_out='NESTED')
+        distmean, diststd = parameters_to_marginal_moments(combined_prob,
+                                                           dist_mu,
+                                                           dist_sigma)
+        out_data = (combined_prob, dist_mu, dist_sigma, dist_norm)
+        out_kwargs['distmean'] = distmean
+        out_kwargs['diststd'] = diststd
+    else:
+        out_data = combined_prob
+
+    # save input headers in output history
+    out_kwargs['HISTORY'] = []
+    for i, x in enumerate(input_skymaps):
+        out_kwargs['HISTORY'].append('')
+        out_kwargs['HISTORY'].append('Headers of HDUs 0 and 1 of input file {:d}:'.format(i))
+        out_kwargs['HISTORY'].append('')
+        out_kwargs['HISTORY'] += ['{} = {}'.format(k, v) for k, v in x[3].items()]
+
+    write_sky_map(args.output, out_data, **out_kwargs)

ligo/skymap/tool/tests/test_tool.py

@@ -3,7 +3,12 @@ import subprocess
 
 import pytest
 
+import numpy as np
+import healpy as hp
+
 from . import run_entry_point
+from ...io import read_sky_map, write_sky_map
+from ...distance import parameters_to_moments
 
 
 @pytest.fixture
@@ -77,3 +82,60 @@ def test_stats(inj_coinc_sqlite, localize_coincs, tmpdir):
     filename = str(tmpdir / 'bayestar.out')
     run_entry_point('ligo-skymap-stats', '-o', filename,
                     inj_coinc_sqlite, os.path.join(localize_coincs, '*.fits'))
+
+
+def test_combine(tmpdir):
+    fn1 = str(tmpdir / 'skymap1.fits.gz')
+    fn2 = str(tmpdir / 'skymap2.fits.gz')
+    fn3 = str(tmpdir / 'joint_skymap.fits.gz')
+
+    # generate a hemisphere of constant probability
+    nside1 = 32
+    npix1 = hp.nside2npix(nside1)
+    m1 = np.zeros(npix1)
+    disc_idx = hp.query_disc(nside1, (1, 0, 0), np.pi/2)
+    m1[disc_idx] = 1
+    m1 /= m1.sum()
+    hp.write_map(fn1, m1, column_names=['PROBABILITY'],
+                 extra_header=[('INSTRUME', 'X1')])
+
+    # generate another hemisphere of constant probability
+    # but with higher resolution and rotated 90 degrees
+    nside2 = 64
+    npix2 = hp.nside2npix(nside2)
+    m2 = np.zeros(npix2)
+    disc_idx = hp.query_disc(nside2, (0, 1, 0), np.pi/2)
+    m2[disc_idx] = 1
+    m2 /= m2.sum()
+    hp.write_map(fn2, m2, column_names=['PROBABILITY'],
+                 extra_header=[('INSTRUME', 'Y1')])
+
+    run_entry_point('ligo-skymap-combine', fn1, fn2, fn3)
+
+    m3 = hp.read_map(fn3, nest=True)
+    npix3 = len(m3)
+    nside3 = hp.npix2nside(npix3)
+    pix_area3 = hp.nside2pixarea(nside3)
+
+    # resolution must match the highest original resolution
+    assert npix3 == npix2
+    # probability must be normalized to 1
+    assert m3.sum() == pytest.approx(1)
+    # support must be ¼ of the sphere
+    tolerance = 10 * hp.nside2pixarea(nside1)
+    assert sum(m3 > 0) * pix_area3 == pytest.approx(np.pi, abs=tolerance)
+
+    # generate a BAYESTAR-like map with mock distance information
+    d_mu = np.zeros_like(m1)
+    d_sigma = np.ones_like(m1)
+    d_norm = np.ones_like(m1)
+    write_sky_map(fn1, np.vstack((m1, d_mu, d_sigma, d_norm)).T)
+
+    run_entry_point('ligo-skymap-combine', fn1, fn2, fn3)
+
+    m3, meta3 = read_sky_map(fn3, nest=True, distances=True)
+
+    # check that marginal distance moments match what was simulated
+    mean, std, _ = parameters_to_moments(d_mu[0], d_sigma[0])
+    assert meta3['distmean'] == pytest.approx(mean)
+    assert meta3['diststd'] == pytest.approx(std)

setup.cfg

@@ -79,6 +79,7 @@ bayestar-localize-coincs = ligo.skymap.tool.bayestar_localize_coincs:main
 bayestar-localize-lvalert = ligo.skymap.tool.bayestar_localize_lvalert:main
 bayestar-realize-coincs = ligo.skymap.tool.bayestar_realize_coincs:main
 bayestar-sample-model-psd = ligo.skymap.tool.bayestar_sample_model_psd:main
+ligo-skymap-combine = ligo.skymap.tool.ligo_skymap_combine:main
 ligo-skymap-contour = ligo.skymap.tool.ligo_skymap_contour:main
 ligo-skymap-from-samples = ligo.skymap.tool.ligo_skymap_from_samples:main
 ligo-skymap-plot = ligo.skymap.tool.ligo_skymap_plot:main