Rework of the skymap utilities

Replaces the outdated (and potentially incorrect) skymap code with code
adapated from ligo.skymap itself. This is incorporated into the
CBCResult object as a method, and generated the fits, pickle and skymap.

bilby/gw/result.py

 from __future__ import division
 
+import json
+import os
+
 import matplotlib.pyplot as plt
+from matplotlib import rcParams
 import numpy as np
-import os
 
 from ..core.result import Result as CoreResult
-from ..core.utils import logger
+from ..core.utils import logger, check_directory_exists_and_if_not_mkdir
 from .utils import (plot_spline_pos,
                     spline_angle_xform)
 
@@ -180,5 +183,182 @@ class CompactBinaryCoalesenceResult(CoreResult):
         fig.savefig(filename, bbox_inches='tight')
         plt.close(fig)
 
+    def plot_skymap(
+            self, maxpts=None, trials=5, jobs=1, enable_multiresolution=True,
+            objid=None, instruments=None, geo=False, dpi=600,
+            transparent=False, colorbar=False, contour=[50, 90],
+            annotate=True, cmap='cylon'):
+        """ Generate a fits file and sky map from a result
+
+        Code adapted from ligo.skymap.tool.ligo_skymap_from_samples and
+        ligo.skymap.tool.plot_skymap
+
+        Parameters
+        ----------
+        maxpts: int
+            Number of samples to use, if None all samples are used
+        trials: int
+            Number of trials at each clustering number
+        jobs: int
+            Number of multiple threads
+        enable_multiresolution: bool
+            Generate a multiresolution HEALPix map (default: True)
+        objid: st
+            Event ID to store in FITS header
+        instruments: str
+            Name of detectors
+        geo: bool
+            Plot in geographic coordinates (lat, lon) instead of RA, Dec
+        dpi: int
+            Resolution of figure in fots per inch
+        transparent: bool
+            Save image with transparent background
+        colorbar: bool
+            Show colorbar
+        contour: list
+            List of contour levels to use
+        annotate: bool
+            Annotate image with details
+        cmap: str
+            Name of the colormap to use
+        """
+
+        from ligo.skymap import io
+        from ligo.skymap.bayestar import rasterize
+        from ligo.skymap import version
+        from astropy.time import Time
+        import numpy as np
+        import os
+        import pickle
+        from ligo.skymap.kde import Clustered2Plus1DSkyKDE, Clustered2DSkyKDE
+        import healpy as hp
+        from ligo.skymap.io import fits
+        from ligo.skymap import plot
+        from ligo.skymap import postprocess
+
+        check_directory_exists_and_if_not_mkdir(self.outdir)
+
+        logger.info('Reading samples for skymap')
+        data = self.posterior
+
+        if maxpts is not None and maxpts < len(data):
+            logger.info('Taking random subsample of chain')
+            data = data.sample(maxpts)
+
+        try:
+            pts = data[['ra', 'dec', 'luminosity_distance']].values
+            cls = Clustered2Plus1DSkyKDE
+            distance = True
+        except KeyError:
+            logger.warning("The results file does not contain luminosity_distance")
+            pts = data[['ra', 'dec']].values
+            cls = Clustered2DSkyKDE
+            distance = False
+
+        skypost = cls(pts, trials=trials, multiprocess=jobs)
+
+        obj_filename = os.path.join(self.outdir, '{}_skypost.obj'.format(self.label))
+        logger.info('Pickling skymap to {}'.format(obj_filename))
+        with open(obj_filename, 'wb') as out:
+            pickle.dump(skypost, out)
+
+        logger.info('Making skymap')
+        hpmap = skypost.as_healpix()
+        if not enable_multiresolution:
+            hpmap = rasterize(hpmap)
+
+        hpmap.meta.update(io.fits.metadata_for_version_module(version))
+        hpmap.meta['creator'] = "bilby"
+        hpmap.meta['origin'] = 'LIGO/Virgo'
+        hpmap.meta['gps_creation_time'] = Time.now().gps
+        hpmap.meta['history'] = ""
+        if objid is not None:
+            hpmap.meta['objid'] = objid
+        if instruments:
+            hpmap.meta['instruments'] = instruments
+        if distance:
+            hpmap.meta['distmean'] = np.mean(data['luminosity_distance'])
+            hpmap.meta['diststd'] = np.std(data['luminosity_distance'])
+
+        try:
+            time = data['geocent_time']
+            hpmap.meta['gps_time'] = time.mean()
+        except KeyError:
+            logger.warning('Cannot determine the event time from geocent_time')
+
+        fits_filename = os.path.join(self.outdir, "{}_skymap.fits".format(self.label))
+        logger.info('Saving skymap fits-file to {}'.format(fits_filename))
+        io.write_sky_map(fits_filename, hpmap, nest=True)
+
+        skymap, metadata = fits.read_sky_map(fits_filename, nest=None)
+        nside = hp.npix2nside(len(skymap))
+
+        # Convert sky map from probability to probability per square degree.
+        deg2perpix = hp.nside2pixarea(nside, degrees=True)
+        probperdeg2 = skymap / deg2perpix
+
+        if geo:
+            obstime = Time(metadata['gps_time'], format='gps').utc.isot
+            ax = plt.axes(projection='geo degrees mollweide', obstime=obstime)
+        else:
+            ax = plt.axes(projection='astro hours mollweide')
+        ax.grid()
+
+        # Plot sky map.
+        vmax = probperdeg2.max()
+        img = ax.imshow_hpx(
+            (probperdeg2, 'ICRS'), nested=metadata['nest'], vmin=0., vmax=vmax,
+            cmap=cmap)
+
+        # Add colorbar.
+        if colorbar:
+            cb = plot.colorbar(img)
+            cb.set_label(r'prob. per deg$^2$')
+
+        if contour is not None:
+            cls = 100 * postprocess.find_greedy_credible_levels(skymap)
+            cs = ax.contour_hpx(
+                (cls, 'ICRS'), nested=metadata['nest'],
+                colors='k', linewidths=0.5, levels=contour)
+            fmt = r'%g\%%' if rcParams['text.usetex'] else '%g%%'
+            plt.clabel(cs, fmt=fmt, fontsize=6, inline=True)
+
+        # Add continents.
+        if geo:
+            geojson_filename = os.path.join(
+                os.path.dirname(plot.__file__), 'ne_simplified_coastline.json')
+            with open(geojson_filename, 'r') as geojson_file:
+                geoms = json.load(geojson_file)['geometries']
+            verts = [coord for geom in geoms
+                     for coord in zip(*geom['coordinates'])]
+            plt.plot(*verts, color='0.5', linewidth=0.5,
+                     transform=ax.get_transform('world'))
+
+        # Add a white outline to all text to make it stand out from the background.
+        plot.outline_text(ax)
+
+        if annotate:
+            text = []
+            try:
+                objid = metadata['objid']
+            except KeyError:
+                pass
+            else:
+                text.append('event ID: {}'.format(objid))
+            if contour:
+                pp = np.round(contour).astype(int)
+                ii = np.round(np.searchsorted(np.sort(cls), contour) *
+                              deg2perpix).astype(int)
+                for i, p in zip(ii, pp):
+                    # FIXME: use Unicode symbol instead of TeX '$^2$'
+                    # because of broken fonts on Scientific Linux 7.
+                    text.append(
+                        u'{:d}% area: {:d} deg²'.format(p, i, grouping=True))
+            ax.text(1, 1, '\n'.join(text), transform=ax.transAxes, ha='right')
+
+        filename = os.path.join(self.outdir, "{}_skymap.png".format(self.label))
+        logger.info("Generating 2D projected skymap to {}".format(filename))
+        plt.savefig(filename, dpi=500)
+
 
 CBCResult = CompactBinaryCoalesenceResult

bilby/gw/utils.py

@@ -568,84 +568,6 @@ def gw_data_find(observatory, gps_start_time, duration, calibration,
     return output_cache_file
 
 
-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))
-
-
 def build_roq_weights(data, basis, deltaF):
 
     """