#!/usr/bin/env python # coding: utf-8 import warnings warnings.simplefilter("ignore") #Used to suppress the RuntimeWarnings. But if you are changing the code for TESTING, MAKE SURE TO COMMENT IT OUT. ''' This script calculates the RDGW150914 constraints with ptemcee - specifically, the n=1 case varying the fundamental frequencies, tshift = 0. It produces the chain plot, corner plot, parameter constraints, and data plotted with the 1-sigma band. Since we are working specifically for the n=1 case, we also compare alpha_0 and alpha_1 in a histogram, in order to demonstrate that the two are indistinguishable with each other. --- R^a_{yne} L^i_u, 08/09/2020 Adapted to include the nmax = 0 case. ---R^a_{yne} L^i_u, 08/20/2020, don't know what to do now ''' import numpy as np import corner import matplotlib.pyplot as plt from matplotlib.ticker import MaxNLocator from matplotlib import rc #plt.rcParams['font.family'] = 'DejaVu Sans' #rc('text', usetex=True) plt.rcParams.update({'font.size': 16.5}) import ptemcee from pycbc.pool import choose_pool import h5py import inspect import pandas as pd import json import qnm import random #Remember to change the following global variables #rootpath: root path to nr data #npoints: number of points you re using for your sampling #nmax: tone index --> nmax = 0 if fitting the fundamental tone #tshift: time shift after the strain peak #vary_fund: whether you vary the fundamental frequency. Works in the model_dv function. rootpath= "/Users/RayneLiu"#"/work/rayne.liu" nmax=0 tshift=0 vary_fund = True #sampler parameters npoints=2400 nwalkers = 1200 ntemps=16 dim = nmax+1 ndim = 4*dim burnin = 1200 #How many points do you burn before doing the corner plot. You need to watch the convergence of the chain plot a bit. #This is trivial but often forgotten: this cannot be more than npoints! Usually 1/5~1/4 npoints is what I observe. numbins = 42 #corner plot parameter - how many bins you want datacolor = '#105670' #'#4fa3a7' pkcolor = '#f2c977' #'#ffb45f' mediancolor = '#f7695c' #'#9b2814' #Import data and necessary functions #TimeOfMaximum def FindTmaximum(y): #Determines the maximum absolute value of the complex waveform absval = y[:,1]*y[:,1]+y[:,2]*y[:,2] vmax=np.max(absval) index = np.argmax(absval == vmax) timemax=gw_sxs_bbh_0305[index,0] return timemax #This loads the 22 mode data gw = {} gw["SXS:BBH:0305"] = h5py.File(rootpath+"/git/rdstackingproject/SXS/BBH_SKS_d14.3_q1.22_sA_0_0_0.330_sB_0_0_-0.440/Lev6/rhOverM_Asymptotic_GeometricUnits_CoM.h5", 'r') gw_sxs_bbh_0305 = gw["SXS:BBH:0305"]["Extrapolated_N2.dir"]["Y_l2_m2.dat"] # Remember to download metadata.json from the simulation with number: 0305. Download Lev6/metadata.json # This postprocesses the metadata file to find the final mass and final spin metadata = {} with open(rootpath+"/git/rdstackingproject/SXS/BBH_SKS_d14.3_q1.22_sA_0_0_0.330_sB_0_0_-0.440/Lev6/metadata.json") as file: metadata["SXS:BBH:0305"] = json.load(file) af = metadata["SXS:BBH:0305"]['remnant_dimensionless_spin'][-1] mf = metadata["SXS:BBH:0305"]['remnant_mass'] #times --> x axis of your data times = gw_sxs_bbh_0305[:,0] tmax=FindTmaximum(gw_sxs_bbh_0305) t0=tmax +tshift #Select the data from t0 onwards position = np.argmax(times >= (t0)) gw_sxs_bbh_0305rd=gw_sxs_bbh_0305[position:-1] timesrd=gw_sxs_bbh_0305[position:-1][:,0] # Depending on nmax, you load nmax number of freqs. and damping times from the qnm package omegas = [qnm.modes_cache(s=-2,l=2,m=2,n=i)(a=af)[0] for i in range (0,dim)] w = (np.real(omegas))/mf tau=-1/(np.imag(omegas))*mf #Fitting #RD model for nmax tones. Amplitudes are in (xn*Exp[i yn]) version. Used here. def model_dv(theta): #x0, y0= theta #Your nmax might not align with the dim of theta. Better check it here. assert int(len(theta)/4) == dim, 'Please recheck your n and parameters' avars = theta[ : (dim)] bvars = theta[(dim) : 2*(dim)] xvars = theta[2*(dim) : 3*(dim)] yvars = theta[3*(dim) : ] if vary_fund == False: avars[0]=0 bvars[0]=0 ansatz = 0 for i in range (0,dim): #bvars[1]=0 #avars[1]=0 ansatz += (xvars[i]*np.exp(1j*yvars[i]))*np.exp(-(timesrd-timesrd[0])/(tau[i]*(1+bvars[i]))) * (np.cos((1+avars[i])*w[i]*timesrd)-1j*np.sin((1+avars[i])*w[i]*timesrd)) # -1j to agree with SXS convention return ansatz # Logprior distribution. It defines the allowed range my variables can vary over. #It works for the (xn*Exp[iyn]) version. def log_prior(theta): #Warning: we are specifically working with nmax=1 so here individual prior to the parameters are manually adjusted. This does not apply to all other nmax's. avars = theta[ : (dim)] bvars = theta[(dim) : 2*(dim)] xvars = theta[2*(dim) : 3*(dim)] yvars = theta[3*(dim) : ] if nmax == 0: if all([0 <= tshift <= 5, vary_fund == True, -0.45 <= avars[0] <= 0.05, -0.95 <= bvars[0] <= 0.95, 0 <= xvars[0] <= 2.4, -np.pi <= yvars[0] <= np.pi]): return 0.0 elif nmax == 1: if all([0 <= tshift <= 5, vary_fund == True, -1.2 <= avars[0] <= 1.2, -1.5 <= avars[1] <= 1.5, -4.0 <= bvars[0] <= 9.0, -4.0 <= bvars[1] <= 21.0, 0 <= xvars[0] <= 1.6, 0 <= xvars[1] <= 1.4, 0 <= yvars[0] <= 2*np.pi, 0 <= yvars[1] <= 2*np.pi]): return 0.0 elif all([tshift == 19, vary_fund == True, -10.0 <= avars[0] <= 10.0, -10.0 <= avars[1] <= 10.0, -20.0 <= bvars[0] <= 30.0, -25.0 <= bvars[1] <= 30.0, 0 <= xvars[0] <= 0.6, 0 <= xvars[1] <= 0.9, 0 <= yvars[0] <= 2*np.pi, 0 <= yvars[1] <= 2*np.pi]): return 0.0 #PAY EXTRA ATTENTION TO THESE TWO CASES, SINCE WE SHIFTED y's. THE CORNER PLOTS LABELS NEED TO BE TREATED WITH CARE. elif all([0 <= tshift <= 5, vary_fund == False, -10.0 <= avars[0] <= 10.0, -1.5 <= avars[1] <= 1.5, -9.0 <= bvars[0] <= 9.0, -6.0 <= bvars[1] <= 20.0, 0 <= xvars[0] <= 2.4, 0 <= xvars[1] <= 2.5, 0 <= yvars[0]-np.pi <= 2*np.pi, 0 <= yvars[1]-np.pi <= 2*np.pi,]): return 0.0 elif all([tshift == 19, vary_fund == False, -10.0 <= avars[0] <= 10.0, -8.0 <= avars[1] <= 8.0, -9.0 <= bvars[0] <= 9.0, -10.0 <= bvars[1] <= 42.0, 0 <= xvars[0] <= 0.6, 0 <= xvars[1] <= 0.7, 0 <= yvars[0]-np.pi <= 2*np.pi, 0 <= yvars[1]-np.pi <= 2*np.pi]): return 0.0 return -np.inf # LogLikelihood function. It is just a Gaussian loglikelihood based on computing the residuals^2 def log_likelihood(theta): modelev = model_dv(theta) result = -np.sum((gw_sxs_bbh_0305rd[:,1] - (modelev.real))**2+(gw_sxs_bbh_0305rd[:,2] - (modelev.imag))**2) if np.isnan(result): return -np.inf return result # Logposterior distribution for the residuals case. # The evidence is just a normalization factor def log_probability(theta): lp = log_prior(theta) if not np.isfinite(lp): return -np.inf return lp + log_likelihood(theta) #Fit with ptemcee #Set the number of cores of your processors pool = choose_pool(24) pool.size = 24 vary_param = float(vary_fund) np.random.seed(42) pos = np.array([[random.uniform(-0.1,0.1), random.uniform(-0.1,0.1), 4.28313743e-01, random.uniform(0.5, 0.6) + (1-vary_param) * np.pi]]) for i in range (1,dim): pos_aux = np.array([[random.uniform(-0.1,0.1), random.uniform(-0.1,0.1), random.uniform(0.3,0.4), random.uniform(0.01, 0.02) + (1-vary_param) * np.pi]]) pos = np.concatenate((pos, pos_aux), axis = 0) pos = pos.T.flatten() pos = list(pos) pos += 1e-5 * np.random.randn(ntemps, nwalkers, ndim) #print(pos) sampler = ptemcee.Sampler(nwalkers, ndim, log_likelihood, log_prior, ntemps=ntemps, pool=pool) sampler.run_mcmc(pos,npoints) #Define labels and start plotting paramlabels_a = [r'$\alpha_'+str(i)+'$' for i in range (dim)] paramlabels_b = [r'$\beta_'+str(i)+'$' for i in range (dim)] paramlabels_x = [r'$x_'+str(i)+'$' for i in range (dim)] paramlabels_y = [r'$y_'+str(i)+'$' for i in range (dim)] if vary_fund == True else ['$y_'+str(i)+'-\pi$' for i in range (dim)] paramlabels = paramlabels_a + paramlabels_b + paramlabels_x + paramlabels_y #Need to delete alpha_0 and alpha_1 for the corner plot paramlabels_corner = paramlabels_a + paramlabels_b + paramlabels_x + paramlabels_y if vary_fund == False: del paramlabels_corner[0] del paramlabels_corner[1] #Chain plot fig, axes = plt.subplots(ndim, 1, sharex=True, figsize=(12, 9*(dim))) for i in range(ndim): axes[i].plot(sampler.chain[0,:, :, i].T, color="k", alpha=0.4, rasterized=True) axes[i].yaxis.set_major_locator(MaxNLocator(5)) axes[i].set_ylabel(paramlabels[i]) axes[-1].set_xlabel('Iterations') #plt.show() fig.savefig(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_chain.png', format = 'png', bbox_inches = 'tight') #Burn samples, calculate peak likelihood value (not necessarily so in atlas) and make corner plot samples = sampler.chain[0,:, burnin:, :].reshape((-1, ndim)) #samples for corner plot samples_corn = samples if vary_fund == True else np.delete(samples, np.s_[0,2], 1) #print('Values with peak likelihood:') lglk = np.array([log_likelihood(samples[i]) for i in range(len(samples))]) pk = samples[np.argmax(lglk)] #print('pk:') #print(pk) pk_corn = pk if vary_fund == True else np.delete(pk, [0,2]) #y_0 range needs some messaging to make the plot. But in order to make the whole picture consistent, better change the range of y_1 too. if vary_fund == False: samples_corn.T[-dim:] -= np.pi #This indeed changes samples_corn itself pk[-dim:] -= np.pi #print('pkFalse:') #print(pk) #print(pk) #Now calculate median (50-percentile) value median = np.median(samples_corn, axis=0) #print(samples) #print(samples_corn) figcorn = corner.corner(samples_corn, bins = numbins, hist_bin_factor = 5, color = datacolor, truths=pk_corn, truth_color = pkcolor, plot_contours = True, labels = paramlabels_corner, quantiles=(0.05, 0.16, 0.5, 0.84, 0.95), levels=[1-np.exp(-0.5), 1-np.exp(-1.64 ** 2/2)], show_titles=True) #Extract the axes in order to add more important line plots naxes = len(pk_corn) axes = np.array(figcorn.axes).reshape((naxes, naxes)) # Loop over the diagonal for i in range(naxes): ax = axes[i, i] ax.axvline(median[i], color=mediancolor) # Loop over the histograms for yi in range(naxes): for xi in range(yi): ax = axes[yi, xi] ax.axvline(median[xi], color=mediancolor) ax.axhline(median[yi], color=mediancolor) ax.plot(median[xi], median[yi], color = mediancolor, marker = 's') figcorn.savefig(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_corner.pdf', format = 'pdf') #Now export the covariance and correlation matrices, as well as the logevidence and autocorrelation time #The covariant matrix samplesT = samples_corn.T cov_matr=np.cov(samplesT) #print('Covariant matrix') df1 = pd.DataFrame(cov_matr, index = paramlabels_corner, columns=paramlabels_corner) df1.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_covariance.h5', key='df') #The correlation matrix corr_matr=np.corrcoef(samplesT) #print('Correlation matrix') df2 = pd.DataFrame(corr_matr,index = paramlabels_corner, columns=paramlabels_corner) df2.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_correlation.h5', key='df') #The logevidence (with error) and save it logevidence = sampler.log_evidence_estimate(fburnin=0.2) df3 = pd.DataFrame(logevidence,index = [r'logevidence', r'error'], columns=[r'vary='+str(vary_fund)+', nmax='+str(nmax)+', tshift='+str(tshift)+', '+str(npoints)+'pts']) df3.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_logevidence.h5', key='df') #The autocorrelation time, this time I don't intend to delete alpha0 and beta0, just keep the nan's there since we don't really use this data, only look at it. It's good to keep it for sanity check's sake. atcrr = sampler.get_autocorr_time()[0] df4 = pd.DataFrame(atcrr,index = paramlabels, columns=[r'Autocorrelation time, vary='+str(vary_fund)+', nmax='+str(nmax)+', tshift='+str(tshift)+', '+str(npoints)+'pts']) df4.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_autocrrtime.h5', key='df') #Now plot alpha_0 on top of alpha_1 - only on top, not stacked, and only valid for vary_fund == True if nmax == 1 and vary_fund == True: samplea0 = samplesT[0] samplea1 = samplesT[1] fighist1 = plt.figure(figsize = (8, 6)) n0, bins0, patches0 = plt.hist(samplea0, bins = numbins * 6, alpha = 0.5, label = r'$\alpha_0$') n1, bins1, patches1 = plt.hist(samplea1, bins = numbins * 10, alpha = 0.5, label = r'$\alpha_1$') #n01, bins01, patches01 = plt.hist([samplea0, samplea1], numbins * 10, rwidth = 3, alpha = 0.5, label = [r'$\alpha_0$', r'$\alpha_1$']) plt.legend() fighist1.savefig(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_histtop.pdf', format = 'pdf') #This plot is stacked fighist1 = plt.figure(figsize = (8, 6)) sn01, sbins01, spatches01 = plt.hist([samplea0, samplea1], numbins * 10, rwidth = 1, stacked = True, alpha = 0.5, label = [r'$\alpha_0$', r'$\alpha_1$']) plt.legend() fighist1.savefig(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_histstack.pdf', format = 'pdf') #Now plot the NR data against the mcmc fit data, together with the 1-sigma varying error data onesig_bounds = np.array([np.percentile(samples[:, i], [16, 84]) for i in range(len(samples[0]))]).T modelfitpk = model_dv(pk) figband = plt.figure(figsize = (12, 9)) #Plot the 1-sigma_percentile for j in range(len(samples)): sample = samples[j] if np.all(onesig_bounds[0] <= sample) and np.all(sample <= onesig_bounds[1]): plt.plot(timesrd, model_dv(sample).real, "#79CAF2", alpha=0.3) plt.plot(timesrd, gw_sxs_bbh_0305rd[:,1], "k", alpha=0.7, lw=2, label=r'NR_re') plt.plot(timesrd, modelfitpk.real, "r", alpha=0.7, lw=2, label=r'FitMCmax_re') plt.title(r'Comparison of the MC fit data and the $1-\sigma$ error band') plt.legend() plt.xlim(timesrd[0], timesrd[0]+80) plt.xlabel("t") plt.ylabel("h") figband.savefig(rootpath+'/git/rdstackingproject/plotsmc/vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_band.png', format = 'png', dpi = 384, bbox_inches = 'tight') #import os #os.system('afplay /System/Library/Sounds/Submarine.aiff')