Update megasky with ligo.skymap

Fix megasky code calls to work with ligo.skymap API. Printing of sky area not supported for this commit.

• Working example of skymap png
• Working example of fits file which can be injested by ligo-skymap-plot and turned into a skymap figure.

BayesWaveUtils/scripts/megasky.py

@@ -4,15 +4,14 @@ import matplotlib
 matplotlib.use('Agg')
 import numpy as np
 import matplotlib.pyplot as plt
-import glob
 import os
 import sys
 from lal import GreenwichSiderealTime
-from bayeswave_plot import find_trig as ft
 import healpy as hp
-import sky_area.sky_area_clustering as sky
+from ligo.skymap import kde, plot, bayestar, postprocess, io
+from astropy.coordinates import SkyCoord
+from astropy.time import Time
 import getopt
-from bayeswave_plot.readbwb import BwbParams
 import math
 from optparse import OptionParser
 
@@ -28,9 +27,6 @@ class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
 lsctables.use_in(LIGOLWContentHandler)
 
 
-import lalinference.io.fits as lf
-import lalinference.plot as lp
-import lalinference.plot.cmap
 def parser():
     """
     Parser for input (command line and ini file)
@@ -42,21 +38,22 @@ def parser():
     parser.add_option("-N", "--nside", type=int, default=128)
     parser.add_option("-I", "--inj", type=str, default=None, help='Injection xml')
     parser.add_option("-e", "--eventnum", type=int, default=None, help='Injection event id')
-    parser.add_option("-r", "--ra", type=float, default=None)
-    parser.add_option("-d", "--dec", type=float, default=None)
-    parser.add_option("--geo", default=False, action="store_true")
+    parser.add_option("-r", "--ra", type=float, default=None, help="Right Ascension of injection")
+    parser.add_option("-d", "--dec", type=float, default=None, help="Declination of injection")
+    parser.add_option("-n", "--npost", type=int, default=5000, help="Number of points to use in making skymap")
+    parser.add_option("--geo", default=False, action="store_true", help="Whether to project onto world map")
     
     (opts,args) = parser.parse_args()
 
     if len(args)==0:
-        print >> sys.stderr, "ERROR: require BW directory"
+        print("ERROR: require BW directory", file=sys.stderr)
         sys.exit()
     if not os.path.isdir(args[0]):
-        print >> sys.stderr, "ERROR: BW dir %s does not exist"%args[0]
+        print("ERROR: BW dir %s does not exist"%args[0], file=sys.stderr)
         sys.exit()
 
 
-    return opts, args,
+    return opts, args
 
 def gps_to_mjd(gps_time):
     """
@@ -140,7 +137,7 @@ def gps_to_iso8601(gps_time):
 # Works with single output directory from 
 # BayesWaveBurst
 # -------------------------------------------
-def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
+def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000, geo=False):
 
     # -- Close any open figures
     plt.close('all')
@@ -148,14 +145,14 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
     # Read S6 MDC log file
     # -----------------
     if mdc is None:
-        print "No MDC log provided."
+        print("No MDC log provided.")
 
     # -- Save PWD for future reference
     topdir = os.getcwd()
 
     # -- Enter requested directory
     os.chdir(directory)
-    print "Entered", os.getcwd()
+    print("Entered", os.getcwd())
 
     try:
         jobname = 'bayeswave.run'
@@ -167,7 +164,6 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
     # Extract information from bayeswave.run
     # Note: We may only need the trigger time from this
     # ---------------------------------------
-#bayeswave = open(jobname, 'r')
     ifoNames = []
     for line in bayeswave:
         spl = line.split()
@@ -197,18 +193,18 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
     # --------------------------------
 
     # -- Input skymap data
-    print "Extracting RA/DEC samples"
-    filename = './chains/' + 'signal_params_h0.dat.0'
+    print("Extracting RA/DEC samples")
+    filename = './chains/' + 'signal_params.dat.0'
     data = np.loadtxt(filename, unpack=True,usecols=(0,1,2))
     ralist = data[1]
     sin_dec = data[2]
-    print "Total samples are {0}".format(ralist.size)
+    print("Total samples are {0}".format(ralist.size))
 
     # -- Remove burn in samples
-    burnin = ralist.size/4
+    burnin = int(ralist.size/4)
     ralist = ralist[burnin:]
     sin_dec = sin_dec[burnin:]
-    print "After removing burn-in samples are {0}".format(ralist.size)
+    print("After removing burn-in samples are {0}".format(ralist.size))
 
     declist = np.arcsin(sin_dec)
     thetalist = np.pi/2.0 - declist
@@ -218,19 +214,19 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
     if radec.shape[0] > npost:
         radec = np.random.permutation(radec)[:npost,:]
 
-    kde = sky.ClusteredSkyKDEPosterior(radec)
+    skypost = kde.Clustered2DSkyKDE(radec)
 
     # -- Get the sidereal time
-    print "Finding sidereal time"
-    print "Using GPS {0}".format(gps)
+    print("Finding sidereal time")
+    print("Using GPS {0}".format(gps))
     trigtime = float(gps)
     lalgps = lal.LIGOTimeGPS(trigtime)
     sidtime = GreenwichSiderealTime(lalgps, 0)
     sidtime = sidtime % (np.pi*2)
 
     # -- Get the injection location
-    print "GOT MDC?"
-    print mdc
+    print("GOT MDC?")
+    print(mdc)
     if mdc is None:
         injtheta = 0
         injphi   = 0
@@ -240,8 +236,8 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
         injtheta, injphi = mdc.get_theta_phi(trigtime)
         injra = injphi + sidtime
         injdec = np.pi/2 - injtheta
-        print "GOT INJECTION PARAMTERS"
-        print injtheta, injra
+        print("GOT INJECTION PARAMTERS")
+        print(injtheta, injra)
 
     # -- Make plots directory, if needed
     plotsDir = './plots'
@@ -250,72 +246,78 @@ def make_skyview(directory='.', mdc=None, NSIDE=128, inj=None, npost=5000):
 
  
     # -- Plot the skymap and injection location
-    skymap = kde.as_healpix(NSIDE, nest=True)
+    skymap = skypost.as_healpix()
+    rhmap = bayestar.rasterize(skymap)
+    nside = hp.npix2nside(len(rhmap))
+    deg2perpix = hp.nside2pixarea(nside, degrees=True)
+    probperdeg2 = rhmap['PROB']/deg2perpix
+    vmax = probperdeg2.max()
 
    
     fig = plt.figure(figsize=(8,6), frameon=False)
-    ax = plt.subplot(111, projection='astro hours mollweide')
-    ax.cla()
+    if(geo):
+      obstime = Time(trigtime, format='gps').utc.isot
+      ax = plt.subplot(111, projection='geo degrees mollweide', obstime=obstime)
+    else:
+      ax = plt.subplot(111, projection='astro hours mollweide')
     ax.grid()
-    lp.healpix_heatmap(skymap, nest=True, vmin=0.0, vmax=np.max(skymap), cmap=plt.get_cmap('cylon'))
+    ax.imshow_hpx((probperdeg2, 'ICRS'), nested=True, vmin=0., vmax=vmax, cmap='cylon')
+    cls = 100 * postprocess.find_greedy_credible_levels(rhmap['PROB'].data)
+    cs = ax.contour_hpx((cls, 'ICRS'), nested=True,
+            colors='k', linewidths=0.5, levels=[50,90])
+    fmt = r'%g\%%' if matplotlib.rcParams['text.usetex'] else '%g%%'
+    plt.clabel(cs, fmt=fmt, fontsize=6, inline=True)
+
+    if(geo):
+     import json
+     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'))
 
 
-    # TODO: our .png skymaps don't work anymore. 
     if (inj is not None):
-        plt.plot(inj['ra'], inj['dec'], 'kx', ms=30, mew=1)
+        ax.plot_coord(Skycoord(np.rad2drg(inj['ra']), np.rad2deg(inj['dec']), unit='deg'), '*',markerfacecolor='white', markeredgecolor='black', markersize=10)
     if mdc is not None:
-        plt.plot(injra, injdec, 'kx', ms=30, mew=1)
+        ax.plot_coord(Skycoord(np.rad2drg(injra), np.rad2deg(injdec), unit='deg'), '*', markerfacecolor='white', markeredgecolor='black', markersize=10)
     plt.savefig(plotsDir+'/skymap.png')
     plt.close()
+    
+    rhmap.meta['origin'] = 'LIGO/Virgo'
+    rhmap.meta['gps_creation_time'] = Time.now().gps
+    rhmap.meta['gps_time'] = trigtime
 
-
-
-    # Meg's hacky way to get sky maps to work, because this:  lf.write_sky_map('skymap_{0}.fits'.format(gps), skymap, nest=True, gps_time=trigtime) doesn't work anymore
-    skymap = np.atleast_2d(skymap)
-    extra_header = []
-
-    extra_header.append(('DATE-OBS', gps_to_iso8601(trigtime), 'UTC date of the observation'))
-    extra_header.append(('MJD-OBS', gps_to_mjd(trigtime), 'modified Julian date of the observation'))
-
-    hp.write_map('skymap_{0}.fits'.format(gps), skymap, nest=True, fits_IDL=True, coord='C', column_names=('PROB', 'DISTMU', 'DISTSIGMA', 'DISTNORM')[:len(skymap)], column_units=('pix-1', 'Mpc', 'Mpc', 'Mpc-2')[:len(skymap)], extra_header=extra_header)
-
-
-    # -- Calculate the 50 and 90% credible intervals
-    sq_deg = (180.0/np.pi)**2
-    print "Calculating sky area ..."
-    (area50, area90) = kde.sky_area( [0.5, 0.9] )*sq_deg
-    print "Got area50 of {0}".format(area50)
+    io.write_sky_map('skymap_{0}.fits'.format(gps), rhmap, nest=True)
 
     # -- Get the found contour and searched areas
     if mdc is not None:
-        print "Calculating p value of injection"
-        injcontour = kde.p_values( np.array([[injra, injdec]]) )[0]
-        print "Got contour of {0}".format(injcontour)
+        print("Calculating p value of injection")
+        true_ra = 180/np.pi*injra
+        true_dec = 180/np.pi*injdec
     
-        print "Calculating searched area..."
-        searcharea = (kde.searched_area( np.array([[injra, injdec]]) )*sq_deg)[0]
-        print "Got searched area of {0}".format(searcharea)
-
     elif (inj is not None):
-        print "Calculating p value of command line injection"
-        injcontour = kde.p_values( np.array([inj['ra'], inj['dec']]) )[0]
-        print "Got contour of {0}".format(injcontour)
-
-        print "Calculating searched area..."
-        searcharea = (kde.searched_area( np.array([inj['ra'], inj['dec']]) )*sq_deg)[0]
-        print "Got searched area of {0}".format(searcharea)
-
+        print("Calculating p value of command line injection")
+        true_ra = 180/np.pi*inj['ra']
+        true_dec = 180/np.pi*inj['dec']
     else:
-        injcontour = 0
-        searcharea = 0
-    pixarea = hp.nside2pixarea(NSIDE, degrees=True)
+        true_ra = None
+        true_dec = None
+    (  searched_area, searched_prob, offset, searched_modes, contour_areas, 
+       area_probs, contour_modes, searched_prob_dist, contour_dists, 
+       searched_vol, searched_prob_vol, contour_vols 
+    ) = postprocess.find_injection_moc(skymap,true_ra=true_ra, true_dec=true_dec, contours=[0.5, 0.9])
+
 
-    # -- Make an output file with key statistics
     outfile = open('skystats.txt', 'w')
-    outfile.write("# area50  area90  searcharea  injcontour \n")
-    outfile.write("{0} {1} {2} {3}".format(area50, area90, searcharea, injcontour))
+    outfile.write("# area50  area90  searcharea  p_value \n")
+    outfile.write("{0} {1} {2} {3}".format(contour_areas[0], contour_areas[1], searched_area, searched_prob))
     outfile.close()
 
+    
 
 # -- Write main script for command line running
 if __name__ == "__main__":
@@ -333,9 +335,10 @@ if __name__ == "__main__":
          
     ra = opts.ra
     dec = opts.dec
-    geo = opts.dec
+    geo = opts.geo
     mdc = opts.mdc
     NSIDE = opts.nside
+    npost = opts.npost
 
     injpos = None
     # Checl if injection file is present
@@ -343,9 +346,9 @@ if __name__ == "__main__":
       # If present, check if event id is provided
       # If yes, proceed. If no, throw error message and exit
       if(opts.eventnum is None):
-        print >> sys.stderr, "Provide event num if giving injfile"
+        print("Provide event num if giving injfile", file=sys.stderr)
         sys.exit()
-      print "Loading xml"    
+      print("Loading xml")    
       xmldoc = utils.load_filename(
            opts.inj, contenthandler=LIGOLWContentHandler)
       try:
@@ -368,18 +371,8 @@ if __name__ == "__main__":
     # If ra and dec provided manually, use those
     if ra is not None and dec is not None:
       injpos  ={'ra': ra, 'dec':dec, 'id':0}
-    # If md log provided, use that. Currently,
-    # there is no scheme to assert exclusivity of injection and mdcs
-    if mdc is not None:
-      print 'Reading mdc log {0}'.format(mdc)
-      try:
-        mdc = ft.Mdc(arg)
-      except:
-        print 'Warning! Failed to read mdc'
-        mdc = None
-
     
-    make_skyview(args[0], mdc, NSIDE, injpos)
+    make_skyview(args[0], mdc, NSIDE, injpos, npost, geo)
 
     # Move back to original dir
     os.chdir(topdir)