Adds basic skymap/fits utilities

bilby/gw/utils.py

@@ -405,3 +405,81 @@ def read_frame_file(file_name, start_time, end_time, channel=None, buffer_time=1
     else:
         logger.warning('No data loaded.')
         return None
+
+
+def save_to_fits(posterior, outdir, label):
+    """ Generate a fits file from a posterior array """
+    from astropy.io import fits
+    from astropy.units import pixel
+    from astropy.table import Table
+    import healpy as hp
+    nside = hp.get_nside(posterior)
+    npix = hp.nside2npix(nside)
+    logger.debug('Generating table')
+    m = Table([posterior], names=['PROB'])
+    m['PROB'].unit = pixel ** -1
+
+    ordering = 'RING'
+    extra_header = [('PIXTYPE', 'HEALPIX',
+                     'HEALPIX pixelisation'),
+                    ('ORDERING', ordering,
+                     'Pixel ordering scheme: RING, NESTED, or NUNIQ'),
+                    ('COORDSYS', 'C',
+                     'Ecliptic, Galactic or Celestial (equatorial)'),
+                    ('NSIDE', hp.npix2nside(npix),
+                     'Resolution parameter of HEALPIX'),
+                    ('INDXSCHM', 'IMPLICIT',
+                     'Indexing: IMPLICIT or EXPLICIT')]
+
+    fname = '{}/{}_{}.fits'.format(outdir, label, nside)
+    hdu = fits.table_to_hdu(m)
+    hdu.header.extend(extra_header)
+    hdulist = fits.HDUList([fits.PrimaryHDU(), hdu])
+    logger.debug('Writing to a fits file')
+    hdulist.writeto(fname, overwrite=True)
+
+
+def plot_skymap(result, center='120d -40d', nside=512):
+    """ Generate a sky map from a result """
+    import scipy
+    from astropy.units import deg
+    import healpy as hp
+    import ligo.skymap.plot  # noqa
+    import matplotlib.pyplot as plt
+    logger.debug('Generating skymap')
+
+    logger.debug('Reading in ra and dec, creating kde and converting')
+    ra_dec_radians = result.posterior[['ra', 'dec']].values
+    kde = scipy.stats.gaussian_kde(ra_dec_radians.T)
+    npix = hp.nside2npix(nside)
+    ipix = range(npix)
+    theta, phi = hp.pix2ang(nside, ipix)
+    ra = phi
+    dec = 0.5 * np.pi - theta
+
+    logger.debug('Generating posterior')
+    post = kde(np.row_stack([ra, dec]))
+    post /= np.sum(post * hp.nside2pixarea(nside))
+
+    fig = plt.figure(figsize=(5, 5))
+    ax = plt.axes([0.05, 0.05, 0.9, 0.9],
+                  projection='astro globe',
+                  center=center)
+    ax.coords.grid(True, linestyle='--')
+    lon = ax.coords[0]
+    lat = ax.coords[1]
+    lon.set_ticks(exclude_overlapping=True, spacing=45 * deg)
+    lat.set_ticks(spacing=30 * deg)
+
+    lon.set_major_formatter('dd')
+    lat.set_major_formatter('hh')
+    lon.set_ticklabel(color='k')
+    lat.set_ticklabel(color='k')
+
+    logger.debug('Plotting sky map')
+    ax.imshow_hpx(post)
+
+    lon.set_ticks_visible(False)
+    lat.set_ticks_visible(False)
+
+    fig.savefig('{}/{}_skymap.png'.format(result.outdir, result.label))

examples/injection_examples/plot_skymap.py

+#!/usr/bin/env python
+"""
+Example script which produces posterior samples of ra and dec and generates a
+skymap
+"""
+from __future__ import division, print_function
+import bilby
+
+duration = 4.
+sampling_frequency = 2048.
+outdir = 'outdir'
+label = 'plot_skymap'
+injection_parameters = dict(
+    mass_1=36., mass_2=29., a_1=0.4, a_2=0.3, tilt_1=0.5, tilt_2=1.0,
+    phi_12=1.7, phi_jl=0.3, luminosity_distance=4000., iota=0.4, psi=2.659,
+    phase=1.3, geocent_time=1126259642.413, ra=1.375, dec=-0.2108)
+
+waveform_arguments = dict(waveform_approximant='IMRPhenomPv2',
+                          reference_frequency=50.)
+
+waveform_generator = bilby.gw.WaveformGenerator(
+    duration=duration, sampling_frequency=sampling_frequency,
+    frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
+    parameters=injection_parameters, waveform_arguments=waveform_arguments)
+
+ifos = bilby.gw.detector.InterferometerList(['H1', 'L1'])
+ifos.set_strain_data_from_power_spectral_densities(
+    sampling_frequency=sampling_frequency, duration=duration,
+    start_time=injection_parameters['geocent_time'] - 3)
+ifos.inject_signal(waveform_generator=waveform_generator,
+                   parameters=injection_parameters)
+
+priors = bilby.gw.prior.BBHPriorSet()
+for key in ['a_1', 'a_2', 'tilt_1', 'tilt_2', 'phi_12', 'phi_jl', 'psi',
+            'mass_1', 'mass_2', 'phase', 'geocent_time', 'luminosity_distance',
+            'iota']:
+    priors[key] = injection_parameters[key]
+
+likelihood = bilby.gw.GravitationalWaveTransient(
+    interferometers=ifos, waveform_generator=waveform_generator,
+    time_marginalization=True, phase_marginalization=True,
+    distance_marginalization=False, prior=priors)
+
+result = bilby.run_sampler(
+    likelihood=likelihood, priors=priors, sampler='dynesty', npoints=1000,
+    injection_parameters=injection_parameters, outdir=outdir, label=label)
+
+# make some plots of the outputs
+result.plot_corner()
+
+bilby.gw.utils.plot_skymap(result, nside=2**6)