Enable inclination in skymaps

gwemopt/args.py

@@ -151,9 +151,12 @@ def parse_args(args):
     parser.add_argument("--true_ra", default=30.0, type=float)
     parser.add_argument("--true_dec", default=60.0, type=float)
     parser.add_argument("--true_distance", default=100.0, type=float)
+    parser.add_argument("--true_inclination", default=0.0, type=float)
 
     parser.add_argument("--doTrueLocation", action="store_true", default=False)
     parser.add_argument("--observability_thresh", default=0.05, type=float)
     parser.add_argument("--doObservabilityExit", action="store_true", default=False)
 
+    parser.add_argument("--inclination", action="store_true", default=False)
+
     return parser.parse_args(args=args)

gwemopt/io/skymap.py

@@ -16,6 +16,9 @@
 from ligo.gracedb.rest import GraceDb
 from ligo.skymap.bayestar import rasterize
 from ligo.skymap.io import read_sky_map
+from ligo.skymap.moc import uniq2nest
+from scipy.interpolate import PchipInterpolator
+from scipy.stats import norm
 
 from gwemopt.paths import SKYMAP_DIR
 
@@ -121,6 +124,94 @@ def get_skymap(event_name: str, output_dir: Path = SKYMAP_DIR, rev: int = None)
     return savepath
 
 
+def read_inclination(skymap, params, map_struct):
+    # check if the sky location is input
+    # if not, the maximum posterior point is taken
+    if params["true_location"]:
+        ra = params["true_ra"]
+        dec = params["true_dec"]
+        print(f"Using the input sky location ra={ra}, dec={dec}")
+        # convert them back to theta and phi
+        phi = np.deg2rad(ra)
+        theta = 0.5 * np.pi - np.deg2rad(dec)
+        # make use of the maP nside
+        maP_idx = np.argmax(skymap["PROBDENSITY"])
+        order, _ = uniq2nest(skymap[maP_idx]["UNIQ"])
+        nside = hp.order2nside(order)
+        # get the nested idx for the given sky location
+        nest_idx = hp.ang2pix(nside, theta, phi, nest=True)
+        # find the row with the closest nested index
+        nest_idxs = []
+        for row in skymap:
+            order_per_row, nest_idx_per_row = uniq2nest(row["UNIQ"])
+            if order_per_row == order:
+                nest_idxs.append(nest_idx_per_row)
+            else:
+                nest_idxs.append(0)
+        nest_idxs = np.array(nest_idxs)
+        row = skymap[np.argmin(np.absolute(nest_idxs - nest_idx))]
+    else:
+        print("Using the maP point from the fits file input")
+        maP_idx = np.argmax(skymap["PROBDENSITY"])
+        uniq_idx = skymap[maP_idx]["UNIQ"]
+        # convert to nested indexing and find the location of that index
+        order, nest_idx = uniq2nest(uniq_idx)
+        nside = hp.order2nside(order)
+        theta, phi = hp.pix2ang(nside, int(nest_idx), nest=True)
+        # convert theta and phi to ra and dec
+        ra = np.rad2deg(phi)
+        dec = np.rad2deg(0.5 * np.pi - theta)
+        print(f"The maP location is ra={ra}, dec={dec}")
+        # fetching the skymap row
+        row = skymap[maP_idx]
+
+    # construct the iota prior
+    cosiota_nodes_num = 10
+    cosiota_nodes = np.cos(np.linspace(0, np.pi, cosiota_nodes_num))
+    colnames = ["PROBDENSITY", "DISTMU", "DISTSIGMA", "DISTNORM"]
+    # do an all-in-one interpolation
+    prob_density, dist_mu, dist_sigma, dist_norm = (
+        PchipInterpolator(
+            cosiota_nodes[::-1],
+            row["{}_SAMPLES".format(colname)][::-1],
+        )
+        for colname in colnames
+    )
+    # now have the joint distribution evaluated
+    u = np.linspace(-1, 1, 1000)  # this is cosiota
+    # fetch the fixed distance
+    dL = params["true_distance"]
+    prob_u = (
+        prob_density(u)
+        * dist_norm(u)
+        * np.square(dL)
+        * norm(dist_mu(u), dist_sigma(u)).pdf(dL)
+    )
+
+    iota = np.arccos(u)
+    prob_iota = prob_u * np.absolute(np.sin(iota))
+    # in GW, iota in [0, pi], but in EM, iota in [0, pi/2]
+    # therefore, we need to do a folding
+    # split the domain in half
+    iota_lt_pi2 = iota[iota < np.pi / 2]
+    prob_lt_pi2, prob_gt_pi2 = prob_iota[iota < np.pi / 2], prob_iota[iota >= np.pi / 2]
+    iota_EM = iota_lt_pi2
+    prob_iota_EM = prob_lt_pi2 + prob_gt_pi2[::-1]
+
+    # normalize
+    prob_iota /= np.trapz(iota_EM, prob_iota_EM)
+
+    map_struct["iota_EM"] = iota_EM
+    map_struct["prob_iota_EM"] = prob_iota_EM
+
+    map_struct["prob_density_interp"] = prob_density
+    map_struct["dist_mu_interp"] = dist_mu
+    map_struct["dist_sigma_interp"] = dist_sigma
+    map_struct["dist_norm_interp"] = dist_norm
+
+    return map_struct
+
+
 def read_skymap(params, map_struct=None):
     """
     Read in a skymap and return a map_struct
@@ -178,9 +269,33 @@ def read_skymap(params, map_struct=None):
                         filename, field=(0, 1, 2, 3), verbose=False, h=True
                     )
                 except:
-                    table = read_sky_map(filename, moc=True, distances=True)
+                    skymap = read_sky_map(filename, moc=True, distances=True)
                     order = hp.nside2order(params["nside"])
-                    t = rasterize(table, order)
+
+                    # for colname in skymap.colnames:
+                    #    if colname.startswith('PROB'):
+                    #        newname = colname.replace('PROB', 'PROBDENSITY')
+                    #        skymap.rename_column(colname, newname)
+                    #        skymap[newname] *= len(skymap) / (4 * np.pi)
+                    #        skymap[newname].unit = u.steradian ** -1
+
+                    if "PROBDENSITY_SAMPLES" in skymap.columns:
+                        if params["inclination"]:
+                            map_struct = read_inclination(skymap, params, map_struct)
+
+                        skymap.remove_columns(
+                            [
+                                f"{name}_SAMPLES"
+                                for name in [
+                                    "PROBDENSITY",
+                                    "DISTMU",
+                                    "DISTSIGMA",
+                                    "DISTNORM",
+                                ]
+                            ]
+                        )
+
+                    t = rasterize(skymap)
                     result = t["PROB"], t["DISTMU"], t["DISTSIGMA"], t["DISTNORM"]
                     healpix_data = hp.reorder(result, "NESTED", "RING")
 

gwemopt/params.py

@@ -142,4 +142,8 @@ def params_struct(opts):
     else:
         params["end_time"] = time.Time(opts.end_time, format="isot", scale="utc")
 
+    params["inclination"] = (
+        opts.inclination if hasattr(opts, "true_location") else False
+    )
+
     return params

gwemopt/plotting/__init__.py

@@ -6,6 +6,7 @@
 from gwemopt.plotting.observability import plot_observability
 from gwemopt.plotting.plot_coverage import make_coverage_plots
 from gwemopt.plotting.plot_efficiency import make_efficiency_plots
+from gwemopt.plotting.plot_inclination import plot_inclination
 from gwemopt.plotting.plot_schedule import make_schedule_plots
 from gwemopt.plotting.plot_skymap import plot_skymap
 from gwemopt.plotting.plot_tiles import make_tile_plots

gwemopt/plotting/plot_inclination.py

@@ -0,0 +1,60 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.stats import norm
+
+
+def density_2d(cosine_iota, distance, map_struct):
+    return (
+        map_struct["prob_density_interp"](cosine_iota)
+        * map_struct["dist_norm_interp"](cosine_iota)
+        * distance**2
+        * norm(
+            loc=map_struct["dist_mu_interp"](cosine_iota),
+            scale=map_struct["dist_sigma_interp"](cosine_iota),
+        ).pdf(distance)
+    )
+
+
+def plot_inclination(params, map_struct):
+    """
+    Function to plot the inclination dependence
+    """
+
+    plot_name = params["outputDir"].joinpath("inclination.pdf")
+
+    # we evaluate the probability density on a grid of iota and distance
+    cosine_iota_dense = np.linspace(-1, 1, 1000)
+    distance_dense = np.linspace(0, 250, 1000)
+
+    COS, DIST = np.meshgrid(cosine_iota_dense, distance_dense)
+
+    plt.figure()
+    plt.xlabel("cos(Inclination")
+    plt.ylabel("Distance [Mpc]")
+    plt.contourf(COS, DIST, density_2d(COS, DIST, map_struct), cmap="Blues", levels=100)
+    if params["true_location"]:
+        plt.plot(
+            np.cos(np.deg2rad(params["true_inclination"])),
+            params["true_distance"],
+            "x",
+            color="black",
+            markersize=10,
+        )
+    plt.savefig(plot_name, dpi=200)
+    plt.close()
+
+    plot_name = params["outputDir"].joinpath("inclination_marginal.pdf")
+
+    plt.figure()
+    plt.xlabel("Inclination")
+    plt.ylabel("PDF")
+    plt.plot(map_struct["iota_EM"], map_struct["prob_iota_EM"])
+    if params["true_location"]:
+        if params["true_inclination"] > 90:
+            iota_EM = params["true_inclination"] - 90.0
+        else:
+            iota_EM = params["true_inclination"]
+        plt.axvline(x=np.cos(np.deg2rad(iota_EM)), linestyle="--", color="black")
+
+    plt.savefig(plot_name, dpi=200)
+    plt.close()

gwemopt/run.py

@@ -18,6 +18,7 @@
     make_coverage_plots,
     make_efficiency_plots,
     make_tile_plots,
+    plot_inclination,
     plot_observability,
     plot_skymap,
 )
@@ -67,6 +68,8 @@ def run(args=None):
     if args.doPlots:
         print("Plotting skymap...")
         plot_skymap(params, map_struct)
+        if args.inclination:
+            plot_inclination(params, map_struct)
 
     if args.doObservability:
         print("Calculating observability")