Skip to content
Snippets Groups Projects
Commit 9a212653 authored by Rayne Liu's avatar Rayne Liu
Browse files

0001 Run for mock data

parent 3f2dc687
No related branches found
No related tags found
No related merge requests found
#!/usr/bin/env python
# coding: utf-8
# ### Let's try the 0.0001 stepsize with the n=0 and 1 mock data.
# In[99]:
#Import relevant modules, import data and all that
import numpy as np
from scipy import interpolate
import corner
import os
os.environ['MPLCONFIGDIR'] = '/home/rayne.liu/.config/matplotlib'
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
from multiprocessing import 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= "/work/rayne.liu/git/rdstackingproject"#"/Users/RayneLiu/git/rdstackingproject"
nmax=1
tshift=10
vary_fund = True
#sampler parameters
npoints = 20#1200
nwalkers = 10#256
ntemps=12
dim = nmax+1
ndim = 4*dim
burnin = 10#600 #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! I usually use half the points.
numbins = 32 #corner plot parameter - how many bins you want
datacolor = '#105670' #'#4fa3a7'
pkcolor = '#f2c977' #'#ffb45f'
mediancolor = '#f7695c' #'#9b2814'
#Defines the mock data
t = np.arange(0, 137, 0.1)
w0, tau0, x0, y0 = [0.55578191, 11.74090006, 0.98213669, 1.47067537]
#Can get the w and tau from example nb and amplitude and phase from the 1910 paper
w1, tau1, x1, y1 = [0.54351639, 3.88312743, 4.29386867, -0.79472571]
print('The fundamental tone frequency, damping time, amplitude and phase:')
print(w0, tau0, x0, y0)
print('The n=1 overtone frequency, damping time, amplitude and phase:')
print(w1, tau1, x1, y1)
mockdata = x0*np.exp(1j*y0)*np.exp(-t/(tau0)) * (np.cos(w0*t)-1j*np.sin(w0*t)) + \
x1*np.exp(1j*y1)*np.exp(-t/(tau1)) * (np.cos(w1*t)-1j*np.sin(w1*t))
#We also try a tshift for this one
t0=tshift
#Select the data from t0 onwards
position = np.argmax(t >= (t0))
mockdata_rd=mockdata[position:-1]
timesrd=t[position:-1][:920]
print(t0) #(This checks that timesrd[0] is indeed t0 - acturally this is a bit off due to stepsize issues,
#but nvm, we'll fix it right away)
t0 = timesrd[0]
print(t0)
timespan = timesrd - t0
mockdata_re =mockdata_rd.real[:920]
mockdata_im = mockdata_rd.imag[:920]
# In[84]:
mocknew_re = interpolate.interp1d(timespan, mockdata_re, kind = 'cubic')
mocknew_im = interpolate.interp1d(timespan, mockdata_im, kind = 'cubic')
# In[87]:
timespan_new = np.linspace(0, timespan[-1], len(timespan)*1000)
mockdatanew_re = mocknew_re(timespan_new)
mockdatanew_im = mocknew_im(timespan_new)
# In[92]:
#Test the new interpolated data
figtest = plt.figure(figsize = (12, 8))
plt.plot(timespan, mockdata_re, "r", alpha=0.3, lw=2, label='Before_re')
plt.plot(timespan_new, mockdatanew_re, "b", alpha=0.3, lw=2, label='After_re')
plt.plot(timespan, mockdata_im, alpha=0.3, lw=2, label='Before_im')
plt.plot(timespan_new, mockdatanew_im, alpha=0.3, lw=2, label='After_im')
plt.legend()
figtest.savefig(rootpath + '/plotsmc/0001_10M_mockinterpolated_datatest_wandt.png', format='png', bbox_inches='tight', dpi=300)
# ### Now the interpolation seems nice according to what we have above...let's start sampling!
# In[100]:
#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'
wvars = theta[ : (dim)]
tvars = 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(-timespan_new/tvars[i]) * (np.cos(wvars[i]*timespan_new)-1j*np.sin(wvars[i]*timespan_new))
# -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) : ]
omega0, omega1, tau0, tau1, xvar0, xvar1, yvar0, yvar1 = theta
if tshift == 0:
if all([0.45 <= omega0 <= 0.63, 0.27 <= omega1 <= 0.6, 0. <= tau0 <= 30., 0. <= tau1 <= 20., \
0 <= xvar0 <= 6, 0 <= xvar1 <= 6, -np.pi <= yvar0 <= np.pi, 0. <= yvar1 <= 2*np.pi]):
return 0.0
elif tshift == 10:
if all([0.5 <= omega0 <= 0.6, 0.5 <= omega1 <= 0.6, 9 <= tau0 <= 15, 0. <= tau1 <= 9., \
0. <= xvar0 <= 3., 0. <= xvar1 <= 6., -np.pi <= yvar0 <= np.pi, -np.pi <= yvar1 <= 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((mockdatanew_re - (modelev.real))**2+(mockdatanew_im - (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)
# In[101]:
#This cell uses the tshift=10 results
#Set the number of cores of your processors
#pool = choose_pool(1)
#pool.size = 1
np.random.seed(42)
pos = np.array([random.uniform(0.55,0.57), random.uniform(0.5,0.52), random.uniform(10., 13.7), random.uniform(4.,7.), random.uniform(0.3,0.5), random.uniform(0.3, 0.5), random.uniform(-2., 2.), random.uniform(-2., 2.)])
pos = list(pos)
pos += 1e-5 * np.random.randn(ntemps, nwalkers, ndim)
with Pool() as pool:
sampler = ptemcee.Sampler(nwalkers, ndim, log_likelihood, log_prior, ntemps=ntemps, pool=pool)
sampler.run_mcmc(pos,npoints)
dim = 2
paramlabels_w = [r'$\omega_'+str(i)+'$' for i in range (dim)]
paramlabels_t = [r'$\tau_'+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)]
paramlabels = paramlabels_w + paramlabels_t + paramlabels_x + paramlabels_y
print('The chain plot:')
#Chain plot
figchain, axes = plt.subplots(ndim, 1, sharex=True, figsize=(12, 4*(4)))
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')
figchain.savefig(rootpath + '/plotsmc/0001_10M_mockinterpolated_chainplot_wandt_'+str(nwalkers)+'walkers_'+str(npoints)+'pts.png', format='png', bbox_inches='tight', dpi=300)
for temp in range(ntemps):
dftemp = pd.DataFrame(sampler.chain[temp,:, :, :].reshape((-1, ndim)), columns=paramlabels)
dftemp.to_csv(rootpath+'/plotsmc/0001_10M_mockinterpolated'+'_nmax'+str(nmax)+'_tshift'+str(tshift)+'_'+str(npoints)+'pt_temp'+str(temp)+'_chain.csv', index = False)
print('We\'re using ptemcee. Our constraints:')
#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, 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 + '/plotsmc/0001_10M_mockinterpolated_cornerplot_wandt_'+'nmax'+str(nmax)+'_tshift'+str(tshift)+'_'+str(nwalkers)+'walkers_'+str(npoints)+'pts.png', format='png', bbox_inches='tight', dpi=300)
#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(timespan_new, model_dv(sample).real, "#79CAF2", alpha=0.3)
plt.plot(timespan_new, mockdatanew_re, "k", alpha=0.7, lw=2, label=r'NR_re')
plt.plot(timespan_new, 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.xlabel("t")
plt.ylabel("h")
figband.savefig(rootpath + '/plotsmc/0001_10M_mockinterpolated_waveform_wandt_'+'nmax'+str(nmax)+'_tshift'+str(tshift)+'_'+str(nwalkers)+'walkers_'+str(npoints)+'pts.png', format = 'png', dpi = 384, bbox_inches = 'tight')
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment