diff --git a/src/pop_models/powerlaw_spin_vec/dist_plots.py b/src/pop_models/powerlaw_spin_vec/dist_plots.py index 35024a83e6618506990b67640917e132508c45b2..450cf3f935f1b6700e98c7487afe3e9eceb55a18 100644 --- a/src/pop_models/powerlaw_spin_vec/dist_plots.py +++ b/src/pop_models/powerlaw_spin_vec/dist_plots.py @@ -74,6 +74,30 @@ def _get_args(raw_args): help="File to store plot of effective spin cumulative rate " "distributions.", ) + parser.add_argument( + "plot_ppd_m", + help="File to store plot of component mass PPDs.", + ) + parser.add_argument( + "plot_ppd_M", + help="File to store plot of total mass PPD.", + ) + parser.add_argument( + "plot_ppd_q", + help="File to store plot of mass ratio PPD.", + ) + parser.add_argument( + "plot_ppd_chi", + help="File to store plot of spin magnitude PPDs.", + ) + parser.add_argument( + "plot_ppd_costheta", + help="File to store plot of spin tilt PPDs.", + ) + parser.add_argument( + "plot_ppd_chi_eff", + help="File to store plot of effective spin PPD.", + ) parser.add_argument( "--n-samples", @@ -83,8 +107,8 @@ def _get_args(raw_args): parser.add_argument( "--n-kde-samples", - default=1000, type=int, - help="Number of samples to draw to compute KDEs.", + default=4, type=int, + help="Number of samples to draw from each population to compute KDEs.", ) parser.add_argument( @@ -99,6 +123,13 @@ def _get_args(raw_args): help="Assuming rate is log-uniform, switch to Jefferey's prior.", ) + parser.add_argument( + "--log-scale-n-oom", + type=int, default=6, + help="Number of orders-of-magnitude below peak to show in log-scale " + "plots.", + ) + parser.add_argument( "--seed", type=int, @@ -122,7 +153,7 @@ def _main(raw_args=None): from . import prob from . import utils - from ..utils import empirical_distribution_function + from ..utils import empirical_distribution_function, gaussian_kde if raw_args is None: raw_args = sys.argv[1:] @@ -205,7 +236,9 @@ def _main(raw_args=None): n_posterior = numpy.size(rates) m = numpy.linspace(m_min, m_max, args.n_samples) - chi = numpy.linspace(0.0, 1.0, args.n_samples) + M = numpy.linspace(m_min, M_max, args.n_samples) + q = numpy.linspace(0.0, 1.0, args.n_samples) + chi = q chi_eff = numpy.linspace(-1.0, +1.0, args.n_samples) costheta = chi_eff @@ -222,10 +255,23 @@ def _main(raw_args=None): pcostheta2 = numpy.empty_like(m_pm) R_pcostheta2 = numpy.empty_like(m_pm) - Pchi_eff = numpy.empty_like(m_pm) - R_Pchi_eff = numpy.empty_like(m_pm) + # Pchi_eff = numpy.empty_like(m_pm) + # R_Pchi_eff = numpy.empty_like(m_pm) + m1_samples = numpy.empty( + (n_posterior*args.n_kde_samples,), + dtype=numpy.float64, + ) + m2_samples = numpy.empty_like(m1_samples) + chi1_samples = numpy.empty_like(m1_samples) + chi2_samples = numpy.empty_like(m1_samples) + costheta1_samples = numpy.empty_like(m1_samples) + costheta2_samples = numpy.empty_like(m1_samples) + if args.rate_logU_to_jeffereys: + weight_samples = numpy.empty_like(m1_samples) + else: + weight_samples = None iterables = zip( rates, @@ -266,6 +312,23 @@ def _main(raw_args=None): pcostheta2[i] = prob.marginal_spin_tilt_pdf(costheta, sigma_chi2) R_pcostheta2[i] = rate * pcostheta2[i] + i_samples = slice(i*args.n_kde_samples, (i+1)*args.n_kde_samples) + + ( + m1_samples[i_samples], m2_samples[i_samples], + chi1_samples[i_samples], chi2_samples[i_samples], + costheta1_samples[i_samples], costheta2_samples[i_samples], + ) = prob.joint_rvs( + args.n_kde_samples, + alpha, m_min, m_max, M_max, + alpha_chi1, beta_chi1, alpha_chi2, beta_chi2, + sigma_chi1, sigma_chi2, + rand_state=rand_state, + ).T + + if args.rate_logU_to_jeffereys: + weight_samples[i_samples] = numpy.sqrt(rate) + ''' ( m1_samp, m2_samp, @@ -325,6 +388,19 @@ def _main(raw_args=None): R_Pchi_eff_marginal = numpy.average(R_Pchi_eff, weights=weights, axis=0) ''' + # Posterior predictive distributions + M_tot_samples = m1_samples + m2_samples + q_samples = m2_samples / m1_samples + chi_eff_samples = ( + ( + m1_samples * chi1_samples * costheta1_samples + + m2_samples * chi2_samples * costheta2_samples + ) / M_tot_samples + ) + + + + fig_m_pm, ax_m_pm = plt.subplots() fig_R_pm, ax_R_m_pm = plt.subplots() fig_pchi, ax_pchi = plt.subplots() @@ -335,6 +411,12 @@ def _main(raw_args=None): fig_Pchi_eff, ax_Pchi_eff = plt.subplots() fig_R_Pchi_eff, ax_R_Pchi_eff = plt.subplots() ''' + fig_ppd_m, ax_ppd_m = plt.subplots() + fig_ppd_M, ax_ppd_M = plt.subplots() + fig_ppd_q, ax_ppd_q = plt.subplots() + fig_ppd_chi, ax_ppd_chi = plt.subplots() + fig_ppd_costheta, ax_ppd_costheta = plt.subplots() + fig_ppd_chi_eff, ax_ppd_chi_eff = plt.subplots() figs = [ fig_m_pm, @@ -345,21 +427,57 @@ def _main(raw_args=None): fig_R_pcostheta, # fig_Pchi_eff, # fig_R_Pchi_eff, + fig_ppd_m, + fig_ppd_M, + fig_ppd_q, + fig_ppd_chi, + fig_ppd_costheta, + fig_ppd_chi_eff, ] - for ax in [ax_m_pm, ax_R_m_pm]: - ax.set_xlabel(r"$m_1$ [M$_\odot$]") + axes_ppd = [ + ax_ppd_m, + ax_ppd_M, + ax_ppd_q, + ax_ppd_chi, + ax_ppd_costheta, + ax_ppd_chi_eff, + ] + + # Logarithmic x-axes + ax_logx = [ + ax_m_pm, ax_R_m_pm, + ax_ppd_m, ax_ppd_M, + ] + for ax in ax_logx: ax.set_xscale("log") + + # Logarithmic y-axes + ax_logy = [ + ax_m_pm, ax_R_m_pm, + ax_pchi, ax_R_pchi, + ax_pcostheta, ax_R_pcostheta, + ax_ppd_m, ax_ppd_M, + ] + for ax in ax_logy: ax.set_yscale("log") - for ax in [ax_pchi, ax_R_pchi]: + # Label axes + ax_ppd_m.set_xlabel(r"$m$ [M$_\odot$]") + ax_ppd_M.set_xlabel(r"$M$ [M$_\odot$]") + ax_ppd_q.set_xlabel(r"$q$") + + for ax in [ax_m_pm, ax_R_m_pm]: + ax.set_xlabel(r"$m_1$ [M$_\odot$]") + + for ax in [ax_pchi, ax_R_pchi, ax_ppd_chi]: ax.set_xlabel(r"$\chi$") - ax.set_yscale("log") - for ax in [ax_pcostheta, ax_R_pcostheta]: + for ax in [ax_pcostheta, ax_R_pcostheta, ax_ppd_costheta]: ax.set_xlabel(r"$\cos\theta$") - ax.set_yscale("log") + for ax in [ax_ppd_chi_eff]: + ax.set_xlabel(r"$\chi_{\mathrm{eff}}$") ''' for ax in [ax_Pchi_eff, ax_R_Pchi_eff]: @@ -396,16 +514,38 @@ def _main(raw_args=None): ) ''' + for ax in axes_ppd: + ax.set_ylabel(r"$\mathrm{PPD}$") + plot_percentiles(ax_m_pm, m, m_pm_percentile) + ax_m_pm.set_ylim(utils.limited_oom_range(m_pm, args.log_scale_n_oom)) + plot_percentiles(ax_R_m_pm, m, m_R_pm_percentile) + ax_R_m_pm.set_ylim( + utils.limited_oom_range(m_R_pm, args.log_scale_n_oom) + ) plot_percentiles(ax_pchi, chi, pchi1_percentile, color="#1f77b4") - if not metadata["iid_spins"]: + if metadata["iid_spins"]: + ax_pchi.set_ylim( + utils.limited_oom_range(pchi1, args.log_scale_n_oom) + ) + else: plot_percentiles(ax_pchi, chi, pchi2_percentile, color="#2ca02c") + ax_pchi.set_ylim( + utils.limited_oom_range([pchi1, pchi2], args.log_scale_n_oom) + ) plot_percentiles(ax_R_pchi, chi, R_pchi1_percentile, color="#1f77b4") - if not metadata["iid_spins"]: + if metadata["iid_spins"]: + ax_R_pchi.set_ylim( + utils.limited_oom_range(R_pchi1, args.log_scale_n_oom) + ) + else: plot_percentiles(ax_R_pchi, chi, R_pchi2_percentile, color="#2ca02c") + ax_R_pchi.set_ylim( + utils.limited_oom_range([R_pchi1, R_pchi2], args.log_scale_n_oom) + ) plot_percentiles( ax_pcostheta, costheta, pcostheta1_percentile, @@ -442,6 +582,74 @@ def _main(raw_args=None): ) ''' + # Plot PPDs + pm1 = gaussian_kde( + m1_samples, weights=weight_samples, bw_method="scott", + )(m) + pm2 = gaussian_kde( + m2_samples, weights=weight_samples, bw_method="scott", + )(m) + + ax_ppd_m.plot(m, pm1, color="#1f77b4") + ax_ppd_m.plot(m, pm2, color="#2ca02c") + + ax_ppd_m.set_ylim( + utils.limited_oom_range([pm1, pm2], args.log_scale_n_oom) + ) + + del pm1, pm2 + + pM = gaussian_kde( + M_tot_samples, weights=weight_samples, bw_method="scott", + )(M) + + ax_ppd_M.plot(M, pM, color="black") + + ax_ppd_M.set_ylim( + utils.limited_oom_range([pM], args.log_scale_n_oom) + ) + + del pM + + pq = gaussian_kde( + q_samples, weights=weight_samples, bw_method="scott", + )(q) + ax_ppd_q.plot(q, pq, color="black") + del pq + + pchi1 = gaussian_kde( + chi1_samples, weights=weight_samples, bw_method="scott", + )(chi) + ax_ppd_chi.plot(chi, pchi1, color="#1f77b4") + del pchi1 + + if not metadata["iid_spins"]: + pchi2 = gaussian_kde( + chi2_samples, weights=weight_samples, bw_method="scott", + )(chi) + ax_ppd_chi.plot(chi, pchi2, color="#2ca02c") + del pchi2 + + pcostheta1 = gaussian_kde( + costheta1_samples, weights=weight_samples, bw_method="scott", + )(costheta) + ax_ppd_costheta.plot(costheta, pcostheta1, color="#1f77b4") + del pcostheta1 + + if not metadata["iid_spins"]: + pcostheta2 = gaussian_kde( + costheta2_samples, weights=weight_samples, bw_method="scott", + )(costheta) + ax_ppd_costheta.plot(costheta, pcostheta2, color="#2ca02c") + del pcostheta2 + + pchi_eff = gaussian_kde( + chi_eff_samples, weights=weight_samples, bw_method="scott", + )(chi_eff) + ax_ppd_chi_eff.plot(chi_eff, pchi_eff, color="black") + del pchi_eff + + for ax in [ax_m_pm, ax_R_m_pm]: ax.set_xlim([m_min, M_max]) @@ -468,6 +676,13 @@ def _main(raw_args=None): fig_R_Pchi_eff.savefig(args.plot_R_Pchi_eff) ''' + fig_ppd_m.savefig(args.plot_ppd_m) + fig_ppd_M.savefig(args.plot_ppd_M) + fig_ppd_q.savefig(args.plot_ppd_q) + fig_ppd_chi.savefig(args.plot_ppd_chi) + fig_ppd_costheta.savefig(args.plot_ppd_costheta) + fig_ppd_chi_eff.savefig(args.plot_ppd_chi_eff) + if __name__ == "__main__": import sys diff --git a/src/pop_models/powerlaw_spin_vec/utils.py b/src/pop_models/powerlaw_spin_vec/utils.py index 4ffd662f12845863dc5b04e628cd21dadda4431e..9feca270bed0e8348499551a4164515a2ac8d966 100644 --- a/src/pop_models/powerlaw_spin_vec/utils.py +++ b/src/pop_models/powerlaw_spin_vec/utils.py @@ -92,3 +92,27 @@ def upcast_scalars(arrays): result.append(numpy.tile(arr, shape)) return result + + +def limited_oom_range(ys, n_oom): + """ + Returns the plotting range to use for (a list of) array(s) of y-values + ``ys``, limiting it to a number of orders of magnitude ``n_oom`` from the + peak value. + """ + import numpy + + # Transform to log-scale + log_ys = numpy.log10(ys) + + # Mask any nan's and inf's that will mess this up. + log_ys = numpy.ma.masked_invalid(log_ys) + + # Determine extrema in log-scale and round to integers + log_y_min = numpy.floor(numpy.min(log_ys)) + log_y_max = numpy.ceil(numpy.max(log_ys)) + + # Choose appropriately limited minimum value + log_y_min = max(log_y_min, log_y_max - n_oom) + + return 10**log_y_min, 10**log_y_max