htcondor/divergence.py

+import argparse
+import json
+
+from astropy_healpix import HEALPix
+from astropy.table import join
+from ligo.skymap.io import read_sky_map
+from ligo.skymap.healpix_tree import HEALPIX_MACHINE_ORDER, HEALPIX_MACHINE_NSIDE
+from ligo.skymap.moc import uniq2nest
+import numpy as np
+from scipy.stats import norm
+
+HPX = HEALPix(nside=HEALPIX_MACHINE_NSIDE, order='nested')
+
+
+def plog2p(p):
+    """Calculate p * log2(p), but return 0 if p is 0."""
+    result = np.empty_like(p)
+    zero = p == 0
+    result[zero] = 0
+    p_nonzero = p[~zero]
+    result[~zero] = p_nonzero * np.log2(p_nonzero)
+    return result
+
+
+def js_divergence(p, q, dx):
+    """Calculate the Jenson-Shannon divergence.
+
+    This is the symmetrized Kullback-Leibler (KL) divergence computed using
+    base 2 logarithms and divided by 2. It is symmetric and has a value between
+    0 (for identical distributions) and 1."""
+    integrand = plog2p(p) + plog2p(q) - 2 * plog2p(0.5 * (p + q))
+    return 0.5 * (integrand * dx).sum()
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('skymaps', nargs=2)
+args = parser.parse_args()
+
+# Load sky map, convert from UNIQ indexing to NESTED indexing at machine order
+tables = []
+for arg in args.skymaps:
+    table = read_sky_map(arg, moc=True)
+    order, ipix = uniq2nest(table.columns.pop('UNIQ'))
+    ipix <<= 2 * (HEALPIX_MACHINE_ORDER - order)
+    table['NESTED'] = ipix
+    tables.append(table)
+
+# Resample both sky maps to a common grid
+table = join(
+    *tables, keys='NESTED', join_type='outer', metadata_conflicts='silent')
+table.sort('NESTED')
+for colname, column in table.columns.items():
+    if colname != 'NESTED':
+        for slice in np.ma.clump_masked(column):
+            if slice.start > 0:
+                column[slice] = column[slice.start - 1]
+
+# Calculate area per pixel
+table['AREA'] = np.diff(table['NESTED'], append=HPX.npix) * HPX.pixel_area
+
+# Calculate JS divergence for 2D sky map
+divergence_2d = js_divergence(
+    table['PROBDENSITY_1'],
+    table['PROBDENSITY_2'],
+    table['AREA'])
+
+# Calculate JS divergence for 3D sky map
+max_r = 6 * max(table.meta['distmean'] for table in tables)
+n_r = 1000
+d_r = max_r / n_r
+r = d_r * np.arange(1, n_r)
+p = norm(loc=table['DISTMU_1'], scale=table['DISTSIGMA_1']).pdf(r[:, np.newaxis]) * table['PROBDENSITY_1'] * table['DISTNORM_1']
+q = norm(loc=table['DISTMU_2'], scale=table['DISTSIGMA_2']).pdf(r[:, np.newaxis]) * table['PROBDENSITY_2'] * table['DISTNORM_2']
+# Calculate volume of each voxel, defined as the region within the
+# HEALPix pixel and contained within the two centric spherical shells
+# with radii (r - d_r / 2) and (r + d_r / 2).
+dV = (np.square(r[:, np.newaxis]) + np.square(d_r) / 12) * d_r * table['AREA']
+divergence_3d = js_divergence(p, q, dV)
+
+# Done
+print(json.dumps({'2d': divergence_2d, '3d': divergence_3d}))

htcondor/historical.sh

@@ -5,7 +5,8 @@ mv output/*.fits output/bayestar.fits
 ligo-skymap-plot baseline/bayestar.fits.gz -o baseline/bayestar.png
 ligo-skymap-plot output/bayestar.fits -o output/bayestar.png
 ligo-skymap-plot-volume baseline/bayestar.fits.gz -o baseline/bayestar.volume.png
-ligo-skymap-plot-volume output/bayestar.fits -o output/bayestar.volume.png
+ligo-skymap-plot-volume output/bayestar.fits -o output/bayestar.volume.png --align-to baseline/bayestar.fits.gz
 python -c 'from matplotlib.testing.compare import compare_images; print(compare_images("baseline/bayestar.png", "output/bayestar.png", 0))'
 python -c 'from matplotlib.testing.compare import compare_images; print(compare_images("baseline/bayestar.volume.png", "output/bayestar.volume.png", 0))'
+python ../htcondor/divergence.py baseline/bayestar.fits.gz output/bayestar.fits > output/divergence.json
 fitsheader output/bayestar.fits > output/bayestar.txt