fixed bugs on QNM_Berti function in rdown.py

code_new/Sumit/RD_Fits.py

+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+"""Generate ringdown templates in the time and perform parameter estimation on them.
+"""
+
+
+# In[1]:
+
+
+#Import relevant modules, import data and all that
+import time
+import numpy as np
+import corner
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MaxNLocator
+from matplotlib import rc
+from configparser import ConfigParser
+plt.rcParams.update({'font.size': 16.5})
+
+from multiprocessing import Pool
+import random
+import dynesty
+from dynesty import plotting as dyplot
+from dynesty.utils import resample_equal
+from dynesty import utils as dyfunc
+import os
+import argparse
+import scipy.optimize as optimization
+from scipy.optimize import minimize
+import rdown as rd
+import rdown_pe as rd_pe
+import rdown_utilities as rd_ut
+import read_data as rdata
+
+
+# In[2]:
+
+
+## Loading and running data tested with NR data
+## Loading and running data tested with Mock data
+
+
+# In[3]:
+
+
+# Cell that calls the arguments from your 'config.ini' file. 
+try:
+    parser = argparse.ArgumentParser(description="Simple argument parser")
+    parser.add_argument("-c", action="store", dest="config_file")
+    result = parser.parse_args()
+    config_file=result.config_file
+    parser = ConfigParser()
+    parser.read(config_file)
+    parser.sections()
+except SystemExit: 
+    parser = ConfigParser()
+    parser.read('config_fixed_n6_m_af.ini')
+    parser.sections()
+    pass
+
+
+# In[4]:
+
+
+# Load variables from config file
+(rootpath, simulation_path_1,simulation_path_2, metadata_file , simulation_number, output_folder,
+ export, overwrite, sampler,nr_code, nbcores,tshift,tend,t_align,
+ nmax , npoints, model, error_str, fitnoise, l_int, index_mass,index_spin,
+error_type, error_val, af, mf,tau_var_str,nm_mock)=rdata.read_config_file(parser)
+
+
+# In[5]:
+
+
+# Show configuration options
+dim = nmax+1
+ndim = 4*dim
+numbins = 32 #corner plot parameter - how many bins you want
+    
+print('model:',model)
+print('nmax:',nmax)
+print('nm_mock:',nm_mock)
+print('tshift:',tshift)
+print('error:', error_str)
+print('error value:',error_val)
+print('export:',export)
+print('nr code:',nr_code)
+print('fit noise:',fitnoise)
+
+
+# In[6]:
+
+
+# Create output directories 
+if not os.path.exists(output_folder):
+    os.mkdir(output_folder)
+    print("Directory " , output_folder ,  " Created ")
+
+if nr_code == 'Mock-data':   
+    nm_mock_str = 'rec_with'+parser.get('rd-mock-parameters','nm_mock')+'_'
+else:
+    nm_mock_str=''
+        
+if error_str:
+    output_folder_1=(output_folder+'/'+model+'-nmax'+str(nmax)+'_'+nm_mock_str+str(error_str)+'_'+str(error_type)+'_fitnoise_'+str(fitnoise))
+else:
+    output_folder_1=output_folder+'/'+model+'-nmax'+str(nmax)+'_'+nm_mock_str+str(error_str)+'_'+'fitnoise_'+str(fitnoise)
+
+if not os.path.exists(output_folder_1):
+    os.mkdir(output_folder_1)
+    print("Directory " , output_folder_1 ,  " Created ")
+
+
+# In[7]:
+
+
+# Define output files 
+pars = [simulation_number,model,nmax,tshift,npoints]
+corner_plot = rdata.create_output_files(output_folder_1,pars,'corner_plot')
+corner_plot_extra = rdata.create_output_files(output_folder_1,pars,'corner_plot_extra')
+diagnosis_plot = rdata.create_output_files(output_folder_1,pars,'diagnosis')
+fit_plot = rdata.create_output_files(output_folder_1,pars,'fit')
+samples_file = rdata.create_output_files(output_folder_1,pars,'post_samples')
+results_file = rdata.create_output_files(output_folder_1,pars,'sampler_results')
+sumary_data = rdata.create_output_files(output_folder_1,pars,'log_z')
+best_data = rdata.create_output_files(output_folder_1,pars,'best_vals')
+
+files = [corner_plot,corner_plot_extra,diagnosis_plot,fit_plot,samples_file,results_file,sumary_data,best_data]
+
+
+# In[8]:
+
+
+# Remove old files if overwrite = True
+if overwrite:
+    rd_ut.rm_files(files)
+
+
+# In[9]:
+
+
+#Load NR data, align in time and resize. Plot real part and amplitude. Finally compute the mismatch and the snr estimate
+data = rdata.read_data(nr_code,simulation_path_1,RD=True,tshift=tshift,tend = tend,metadata_file=metadata_file,parser=parser)
+data_l = rdata.read_data(nr_code,simulation_path_2,RD=True,tshift=tshift,tend = tend,metadata_file=metadata_file,parser=parser)
+data_r, data_lr = rdata.nr_resize(data,data_l,tshift=tshift,tend=tend)
+times_rd = data_r[:,0]
+
+plt.figure(figsize = (12, 8))
+plt.plot(times_rd, data_r[:,1].real, "r", alpha=0.3, lw=3, label=r'$Lev6$: real')
+plt.plot(times_rd, np.sqrt((data_r[:,1].real)**2+(data_r[:,1].imag)**2), "r", alpha=0.3, lw=3, label=r'$Lev5\,amp$')
+plt.plot(times_rd, data_lr[:,1].real, "b", alpha=0.3, lw=3, label=r'$Lev5: real$')
+plt.plot(times_rd, np.sqrt((data_lr[:,1].real)**2+(data_lr[:,1].imag)**2), "b", alpha=0.3, lw=3, label=r'$Lev5\,amp$')
+if error_str and error_val==0:
+    error = np.sqrt(data_r[:,1]*data_r[:,1]-2*data_r[:,1]*data_lr[:,1]+data_lr[:,1]*data_lr[:,1])
+    error_est=np.sqrt(error.imag**2+error.real**2)
+    plt.plot(times_rd, error_est, "g", alpha=0.3, lw=2, label='error')
+plt.legend()
+
+mismatch=1-rd_ut.EasyMatchT(times_rd,data_r[:,1],data_lr[:,1],tshift,tend)
+error=np.sqrt(2*mismatch)
+print('error estimate:',error)
+print('mismatch:', mismatch)
+print('snr:', rd_ut.EasySNRT(times_rd,data_r[:,1],data_lr[:,1],tshift,tend)/error**2)
+
+
+# In[10]:
+
+
+# Phase alignement
+if parser.has_option('rd-model','phase_alignment'):
+    phase_alignment=eval(parser.get('rd-model','phase_alignment'))    
+else:
+    phase_alignment=False
+    
+if phase_alignment:
+    datar_al = rdata.phase_align(data_r,data_lr)
+    gwdatanew5 = data_lr[:,1]
+    gwdatanew = datar_al[:,1]    
+    timesrd_final = datar_al[:,0]
+    mismatch=1-rd_ut.EasyMatchT(timesrd_final,gwdatanew,gwdatanew5,tshift,tend)
+    error=np.sqrt(2*mismatch)
+    print('error estimate:',error)
+    print('mismatch:', mismatch)
+    print('snr:', rd_ut.EasySNRT(timesrd_final,gwdatanew,gwdatanew5,tshift,tend)/error)
+    if error_str:
+        error = np.sqrt(gwdatanew*gwdatanew-2*gwdatanew*gwdatanew5+gwdatanew5*gwdatanew5)
+        error_est=np.sqrt(error.imag**2+error.real**2)
+    else :
+        error = 1 
+else:
+    datar_al = data_r
+    timesrd_final = datar_al[:,0]
+    
+#Test the new interpolated data
+if error_str and error_val==0:
+    plt.figure(figsize = (12, 8))
+    plt.plot(timesrd_final, datar_al[:,1].real, "r", alpha=0.3, lw=2, label='Original')
+    plt.plot(timesrd_final, data_lr[:,1].real, "b", alpha=0.3, lw=2, label='Aligned')
+    plt.plot(timesrd_final, error_est, "b", alpha=0.3, lw=2, label='error')
+    plt.legend()
+
+
+# In[11]:
+
+
+# Define your noise depending on the noise configuration. Load priors and setup the likelihood with rd_pe.Ringdown_PE. 
+if error_str and error_val==0:
+    error_final = error_est
+    norm_factor = 100*len(error_final)/2*np.log(2*np.pi)
+elif error_str and error_val!=0:
+    datar_al[:,1]+=random.uniform(0, error_val)
+    datar_al[:,1]+=1j*random.uniform(0, error_val)
+    error_tsh = error_val
+    error_final=(error_tsh.real**2+error_tsh.imag**2)
+    norm_factor = 0
+else:
+    error_tsh=1
+    error_final=(error_tsh.real**2+error_tsh.imag**2)
+    norm_factor = 0
+    
+if parser.has_option('setup','qnm_model'):
+    qnm_model='berti'
+    rdownfolders=np.asarray([rootpath+'/RDmodes/l2/n'+str(i+1)+'l2m2.dat' for i in range(nmax+1)])
+    rdowndata = np.asarray([np.loadtxt(rdownfolders[i]).T for i in range(len(rdownfolders))])
+else:
+    qnm_model='qnm'
+
+priors = rd_pe.load_priors(model,parser,nmax,fitnoise=fitnoise)
+rdown=rd.Ringdown_Spectrum(mf,af,2,2,n=nmax,s=-2,time=timesrd_final, rdowndata = rdowndata)
+rdown_pe = rd_pe.Ringdown_PE(rdown,datar_al,dim,priors,errors2=error_final,norm_factor=norm_factor,model=model,l_int=l_int)
+
+
+# In[13]:
+
+
+# Get a first estimate by trying to fit the data.
+nll = lambda *args: -rdown_pe.log_likelihood(*args)
+if model == 'w-tau-fixed-m-af':
+    if fitnoise:
+        initial = np.concatenate((np.ones(2*dim),[0.8,0.9,1]))
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+    else:
+        initial = np.concatenate((np.ones(2*dim),[0.8,0.9]))
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+elif model == 'w-tau-fixed':
+    if fitnoise:
+        initial =  np.concatenate((np.ones(2*dim),[0.2]))
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+    else:
+        initial = np.ones(2*dim)
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+else:
+    if fitnoise:
+        initial = np.concatenate((np.ones(ndim),[1]))
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+    else: 
+        initial = np.ones(ndim)
+        soln = minimize(nll, initial,bounds=priors)
+        vars_ml=soln.x
+print("best fit pars from fit: ",vars_ml)
+
+
+# In[ ]:
+
+
+mypool = Pool(nbcores)
+mypool.size = nbcores
+
+start = time.process_time()
+f2=dynesty.NestedSampler(rdown_pe.log_likelihood,rdown_pe.prior_transform, len(priors), nlive=npoints,sample=sampler,pool=mypool)
+if parser.has_option('setup','dlogz'):
+    dlogz=np.float(parser.get('setup','dlogz'))    
+    f2.run_nested(dlogz=dlogz,print_progress=False)
+else:
+    f2.run_nested(print_progress=False)
+
+print(time.process_time() - start)
+
+
+# In[52]:
+
+
+res = f2.results
+res.samples_u.shape
+res.summary()
+samps=f2.results.samples
+postsamps = rd_ut.posterior_samples(f2)
+samps_tr=np.transpose(samps)
+half_points=int(round((len(samps_tr[0])/1.25)))
+evidence = res.logz[-1]
+evidence_error = res.logzerr[-1]
+if export:
+    rd_ut.save_object(res, results_file)
+
+
+# In[53]:
+
+
+pars = nmax,model,samps_tr, half_points
+npamps = rd_ut.get_best_amps(pars,parser=parser,nr_code=nr_code)
+
+
+# In[54]:
+
+
+if export:
+    pars = simulation_number, nmax, tshift, evidence, evidence_error
+    rd_ut.export_logz_files(sumary_data,pars)
+
+
+# In[55]:
+
+
+labels = rd_ut.define_labels(dim,model,fitnoise)
+if export:
+    pars = tshift, len(priors), labels 
+    rd_ut.export_bestvals_files(best_data,postsamps,pars)
+
+
+# In[56]:
+
+
+w, tau = rdown.QNM_spectrum()
+pars = w, tau, mf, af, npamps 
+truths = rd_ut.get_truths(model,pars,fitnoise)
+
+
+# In[57]:
+
+
+fg=corner.corner(postsamps,quantiles=[0.05,0.5,0.95],show_titles=True,max_n_ticks = 4,bins=50,truths=truths,labels=labels,truth_color='red')
+plt.show()
+if export:
+    fg.savefig(corner_plot, format = 'png', bbox_inches = 'tight')
+
+
+# In[58]:
+
+
+from importlib import reload
+reload(rd_ut)
+if model == 'w-tau-fixed-m-af' and export == True:        
+    truths=np.concatenate((w,tau))
+    labels_mf = np.concatenate((w_lab,tau_lab))
+    new_samples = rd_ut.convert_m_af_2_w_tau_post(res,fitnoise=False)
+    figure = corner.corner(new_samples,truths=truths,quantiles=[0.05,0.95],labels=labels_mf,smooth=True,color='b',truth_color='r',show_titles=True)
+    figure.savefig(corner_plot_extra, format = 'png', bbox_inches = 'tight')
+
+
+# In[151]:
+
+
+#lnz_truth = ndim * -np.log(2 * 10.)  # analytic evidence solution
+fig, axes = dyplot.runplot(res)
+fig.tight_layout()
+if export:
+    fig.savefig(diagnosis_plot, format = 'png', dpi = 384, bbox_inches = 'tight')
+
+
+# In[166]:
+
+
+if export:
+    dict = {'w-tau':rdown.rd_model_wtau , 'w-q': rdown.rd_model_wq, 'w-tau-fixed':rdown.rd_model_wtau_fixed,'w-tau-fixed-m-af': rdown.rd_model_wtau_m_af}
+    figband = plt.figure(figsize = (12, 9))
+    plt.plot(datar_al[:,0].real,datar_al[:,1].real, "green", alpha=0.9, lw=3, label=r'$res_{240}$')
+    onesig_bounds = np.array([np.percentile(postsamps[:, i], [5, 95]) for i in range(len(postsamps[0]))]).T
+    samples_1sigma = filter(lambda sample: np.all(onesig_bounds[0] <= sample) and np.all(sample <= onesig_bounds[1]), postsamps)
+    samples_1sigma_down = list(samples_1sigma)[::downfactor]
+    for sample in samples_1sigma_down:
+        plt.plot(datar_al[:,0].real, dict[model](sample).real, "r-", alpha=0.1, lw=1)
+    plt.title(r'Comparison of the MC fit data and the $1-\sigma$ error band')
+    plt.legend()
+    plt.xlabel('t')
+    plt.ylabel('h')
+    plt.show()
+    figband.savefig(fit_plot)
+
+
+# In[162]:
+
+
+if export:
+    with open(samples_file,'w') as file:
+        writer = csv.writer(file)
+        writer.writerow(labels)
+        writer.writerows(samps[::downfactor])
+

code_new/Sumit/read_data.py

@@ -226,11 +226,11 @@ def create_output_files(output_folder,pars,file_type):
     """
     
     if file_type=='corner_plot':
-        outfile = output_folder+'/Dynesty_'+str(sim_num)+'_'+model+'_nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'corner_plot.png'
+        outfile = output_folder+'/Dynesty_'+str(sim_num)+'_'+model+'_nmax_'+str(nmax)+'_tshift_'+str(tshift)+'_'+str(npoints)+'corner_plot.png'
     elif file_type=='corner_plot_extra':
-        outfile = output_folder+'/Dynesty_'+str(sim_num)+'_'+model+'_nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'corner_plot_extra.png'
+        outfile = output_folder+'/Dynesty_'+str(sim_num)+'_'+model+'_nmax_'+str(nmax)+'_tshift_'+str(tshift)+'_'+str(npoints)+'corner_plot_extra.png'
     elif file_type=='diagnosis':
-        outfile = output_folder+'/Dynesty_diagnosis'+str(sim_num)+'_'+model+'_nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'.png'
+        outfile = output_folder+'/Dynesty_diagnosis'+str(sim_num)+'_'+model+'_nmax_'+str(nmax)+'_tshift_'+str(tshift)+'_'+str(npoints)+'.png'
     elif file_type=='fit':
         outfile = output_folder+'/Fit_results_'+str(sim_num)+'tshift_'+str(tshift)+'_'+model+'_nmax_'+str(nmax)+'.png'
     elif file_type=='post_samples':
@@ -340,5 +340,5 @@ def read_config_file(parser):
     else:
         nm_mock = None
         
-    res = simulation_path_1,simulation_path_2, metadata_file , simulation_number, output_folder, export, overwrite, sampler,nr_code, nbcores,tshift,tend,t_align, nmax , npoints, model, error_str, fitnoise, l_int, index_mass,index_spin, error_type, error_val, af, mf,tau_var_str,nm_mock
+    res = rootpath, simulation_path_1,simulation_path_2, metadata_file , simulation_number, output_folder, export, overwrite, sampler,nr_code, nbcores,tshift,tend,t_align, nmax , npoints, model, error_str, fitnoise, l_int, index_mass,index_spin, error_type, error_val, af, mf,tau_var_str,nm_mock
     return res
\ No newline at end of file