From 96ec29a1f61f6fac7f0cf0d855e9e3a934f67f41 Mon Sep 17 00:00:00 2001 From: Michael Coughlin Date: Tue, 12 Mar 2024 11:41:51 -0500 Subject: [PATCH] Enable inclination in skymaps --- gwemopt/args.py | 3 + gwemopt/io/skymap.py | 119 ++++++++++++++++++++++++++- gwemopt/params.py | 4 + gwemopt/plotting/__init__.py | 1 + gwemopt/plotting/plot_inclination.py | 60 ++++++++++++++ gwemopt/run.py | 3 + 6 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 gwemopt/plotting/plot_inclination.py diff --git a/gwemopt/args.py b/gwemopt/args.py index 7252b8fc..a5f7316a 100644 --- a/gwemopt/args.py +++ b/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) diff --git a/gwemopt/io/skymap.py b/gwemopt/io/skymap.py index dbd430aa..0ee98623 100644 --- a/gwemopt/io/skymap.py +++ b/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") diff --git a/gwemopt/params.py b/gwemopt/params.py index bcea7ec2..07211a2a 100644 --- a/gwemopt/params.py +++ b/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 diff --git a/gwemopt/plotting/__init__.py b/gwemopt/plotting/__init__.py index 74d4b170..1a87c30f 100644 --- a/gwemopt/plotting/__init__.py +++ b/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 diff --git a/gwemopt/plotting/plot_inclination.py b/gwemopt/plotting/plot_inclination.py new file mode 100644 index 00000000..b85752ef --- /dev/null +++ b/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() diff --git a/gwemopt/run.py b/gwemopt/run.py index ba4968fb..a389be7f 100644 --- a/gwemopt/run.py +++ b/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")