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 pickle
+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 +184,177 @@ 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. Note, the use of this additionally
+        required the installation of ligo.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
+        """
+
+        try:
+            from astropy.time import Time
+            from ligo.skymap import io, version, plot, postprocess, bayestar, kde
+            import healpy as hp
+        except ImportError as e:
+            logger.info("Unable to generate skymap: error {}".format(e))
+            return
+
+        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 = kde.Clustered2Plus1DSkyKDE
+            distance = True
+        except KeyError:
+            logger.warning("The results file does not contain luminosity_distance")
+            pts = data[['ra', 'dec']].values
+            cls = kde.Clustered2DSkyKDE
+            distance = False
+
+        logger.info('Initialising skymap class')
+        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 = bayestar.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 = io.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):
+                    text.append(
+                        u'{:d}% area: {:d} deg$^2$'.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