forked from PoulinV/AxiCLASS
-
Notifications
You must be signed in to change notification settings - Fork 0
/
extract_cs2.py
329 lines (311 loc) · 14.9 KB
/
extract_cs2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
# coding: utf-8
# In[ ]:
# import necessary modules
#get_ipython().magic(u'matplotlib notebook')
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from classy import Class
from scipy.optimize import fsolve
from scipy.interpolate import interp1d
import math
from matplotlib import rc
from scipy.special import gamma, factorial
rc('font',**{'family':'serif','serif':['Times']})
rc('text', usetex=True)
# In[ ]:
# esthetic definitions for the plots
matplotlib.mathtext.rcParams['legend.fontsize']='medium'
plt.rcParams["figure.figsize"] = [8.0,6.0]
# In[ ]:
#############################################
#
# User settings controlling the figure aspect
#
z_max_pk = 46000 # highest redshift involved
k_per_decade = 400 # number of k values, controls final resolution
k_min_tau0 = 40. # this value controls the minimum k value in the figure (it is k_min * tau0)
P_k_max_inv_Mpc =1.0 # this value is directly the maximum k value in the figure in Mpc
tau_num_early = 2000 # number of conformal time values before recombination, controls final resolution
tau_num_late = 200 # number of conformal time values after recombination, controls final resolution
tau_ini = 10. # first value of conformal time in Mpc
tau_label_Hubble = 70. # value of time at which we want to place the label on Hubble crossing
tau_label_ks = 40. # value of time at which we want to place the label on sound horizon crossing
tau_label_kd = 230. # value of time at which we want to place the label on damping scale crossing
#
# Cosmological parameters and other CLASS parameters
#
Contours_at = (0.51,0.7,0.9)
Theta_initial_fld=2.8
N = 10
Theta_initial_fld_array = np.linspace(0.1,3,N)
n_axion = 3
log10ac = -3.5
k9_over_kac=[]
v1 = True
if v1 == False:
Theta_initial_fld_array = np.linspace(0.1,3,N)
else:
Theta_initial_fld_array = [Theta_initial_fld]
for Theta_initial_fld in Theta_initial_fld_array:
common_settings = {'scf_potential' : 'axion',
'output':'mTk',
'n_axion' : n_axion,
'scf_parameters' : '%.1f,0'%(Theta_initial_fld),
'scf_tuning_index':0,
'log10_fraction_axion_ac':-0.8769,
'log10_axion_ac':log10ac,
'omega_cdm':1.320351e-01,
'omega_b': 2.263127e-02,
# '100*theta_s':1.04262e+00,
# '100*theta_s':1.04262e+00,
'H0':7.238797e+01,
'tau_reio':6.402188e-02,
'A_s':2.178373e-09,
'n_s':9.917292e-01,
'scf_evolve_like_axionCAMB':'no',
'do_shooting':'yes',
'do_shooting_scf':'yes',
'use_big_theta_scf':'yes',
'scf_has_perturbations':'yes',
'attractor_ic_scf':'no',
'adptative_stepsize':100,
'scf_evolve_as_fluid':'no',
'z_max_pk':z_max_pk,
'recfast_z_initial':z_max_pk,
#'k_step_sub':'0.01',
'k_per_decade_for_pk':k_per_decade,
'k_per_decade_for_bao':k_per_decade,
'k_min_tau0':k_min_tau0, # this value controls the minimum k value in the figure
'perturb_sampling_stepsize':'0.05',
'P_k_max_1/Mpc':P_k_max_inv_Mpc,
'compute damping scale':'yes', # needed to output and plot Silk damping scale
'gauge':'synchronous'}
# common_settings['h'] = 0.67
# common_settings['back_integration_stepsize'] = 5e-4
###############
#
# call CLASS
#
###############
M = Class()
M.set(common_settings)
M.compute()
#
# define conformal time sampling array
#
times = M.get_current_derived_parameters(['tau_rec','conformal_age'])
tau_rec=times['tau_rec']
tau_0 = times['conformal_age']
tau1 = np.logspace(math.log10(tau_ini),math.log10(tau_rec),tau_num_early)
tau2 = np.logspace(math.log10(tau_rec),math.log10(tau_0),tau_num_late)[1:]
tau2[-1] *= 0.999 # this tiny shift avoids interpolation errors
tau = np.concatenate((tau1,tau2))
tau_num = len(tau)
#
# use table of background and thermodynamics quantitites to define some functions
# returning some characteristic scales
# (of Hubble crossing, sound horizon crossing, etc.) at different time
#
background = M.get_background() # load background table
#print background.viewkeys()
thermodynamics = M.get_thermodynamics() # load thermodynamics table
#print thermodynamics.viewkeys()
#
background_tau = background['conf. time [Mpc]'] # read conformal times in background table
background_z = background['z'] # read redshift
# background_aH = 2.*math.pi*background['H [1/Mpc]']/(1.+background['z'])/M.h() # read 2pi * aH in [h/Mpc]
background_aH = background['H [1/Mpc]']/(1.+background['z'])/M.h() # read 2pi * aH in [h/Mpc]
background_ks = 2.*math.pi/background['comov.snd.hrz.']/M.h() # read 2pi/(comoving sound horizon) in [h/Mpc]
background_rho_m_over_r = (background['(.)rho_b']+background['(.)rho_cdm']) /(background['(.)rho_g']+background['(.)rho_ur']) # read rho_r / rho_m (to find time of equality)
background_rho_l_over_m = background['(.)rho_lambda'] /(background['(.)rho_b']+background['(.)rho_cdm']) # read rho_m / rho_lambda (to find time of equality)
thermodynamics_tau = thermodynamics['conf. time [Mpc]'] # read confromal times in thermodynamics table
# thermodynamics_kd = 2.*math.pi/thermodynamics['r_d']/M.h() # read 2pi(comoving diffusion scale) in [h/Mpc]
#
# define a bunch of interpolation functions based on previous quantities
#
background_z_at_tau = interp1d(background_tau,background_z)
background_tau_at_z = interp1d(background_z,background_tau)
background_aH_at_tau = interp1d(background_tau,background_aH)
background_ks_at_tau = interp1d(background_tau,background_ks)
background_tau_at_mr = interp1d(background_rho_m_over_r,background_tau)
background_tau_at_lm = interp1d(background_rho_l_over_m,background_tau)
#
# infer arrays of characteristic quantitites calculated at values of conformal time in tau array
#
aH = background_aH_at_tau(tau)
ks = background_ks_at_tau(tau)
#
# infer times of R/M and M/Lambda equalities
#
tau_eq = background_tau_at_mr(1.)
tau_lambda = background_tau_at_lm(1.)
#
# check and inform user whether intiial arbitrary choice of z_max_pk was OK
max_z_needed = background_z_at_tau(tau[0])
if max_z_needed > z_max_pk:
print 'you must increase the value of z_max_pk to at least ',max_z_needed
() + 1 # this strange line is just a trick to stop the script execution there
else:
print 'in a next run with the same values of tau, you may decrease z_max_pk from ',z_max_pk,' to ',max_z_needed
#
# get transfer functions at each time and build arrays Theta0(tau,k) and phi(tau,k)
#
i_at_ac =0
for i in range(tau_num):
one_time = M.get_transfer(background_z_at_tau(tau[i])) # transfer functions at each time tau
if i ==0: # if this is the first time in the loop: create the arrays (k, Theta0, phi)
k = one_time['k (h/Mpc)']
k_num = len(k)
Theta0 = np.zeros((tau_num,k_num))
cs2_fld = np.zeros((tau_num,k_num))
phi = np.zeros((tau_num,k_num))
# Theta0[i,:] = 0.25*one_time['d_g'][:]
a = 1/(1+background_z_at_tau(tau[i]))
if a > 10**log10ac and i_at_ac == 0:
i_at_ac=i
a_osc=10**log10ac/1.7
w_n = (n_axion-1)/(n_axion+1)
theta_osc = Theta_initial_fld *(a/a_osc)**(-3*(1+w_n)/(2*n_axion))
c = 3e5
m = 10**M.log10_m_axion()*M.h()*100/c
# print "m:",m
omega = m*np.sqrt(np.pi)*gamma((1.0+n_axion)/(2.0*n_axion))/gamma(1.0+1.0/(2*n_axion))*2**(-(1.0+n_axion)/2.0)*theta_osc**(n_axion-1.0)
# print "omega:",omega
cs2_fld[i,:] = (2*a*a*(n_axion-1)*omega*omega+k*k)/(2*a*a*(n_axion+1)*omega*omega+k*k)
phi[i,:] = one_time['phi'][:]
# Omega_scf = bg['(.)Omega_scf']
H = background['H [1/Mpc]']
Da = background['ang.diam.dist.']
z = background['z']
# tau = background['conf. time [Mpc]']
# H_of_z = interp1d(z,H)
zc = 1/(10**log10ac)-1
tau_ac = background_tau_at_z(zc)
k_of_cs2_at_ac = interp1d(cs2_fld[i_at_ac,:],k)
print cs2_fld[i_at_ac,:],i_at_ac
k_ac = background_aH_at_tau(tau_ac)
print "k_ac=",k_ac,"k_of_cs2_at_ac=",k_of_cs2_at_ac(0.9)
k9_over_kac.append(k_of_cs2_at_ac(0.9)/k_ac)
M.struct_cleanup()
# print "k_ac=",k_ac
# for j in range(len(Omega_scf)):
# if Omega_scf[j] > 0.01:
# k = H[j]/(1+z[j])
#
# find the global extra of Theta0(tau,k) and phi(tau,k), used to define color code later
#
#
#################
#
# start plotting
#
#################
#
if v1 == True:
Theta_amp = max(Theta0.max(),-Theta0.min())
phi_amp = max(phi.max(),-phi.min())
#
# reshaping of (k,tau) necessary to call the function 'pcolormesh'
#
K,T = np.meshgrid(k,tau)
#
# inform user of the size of the grids (related to the figure resolution)
#
print 'grid size:',len(k),len(tau),Theta0.shape
fig = plt.figure(figsize=(12,8))
#
# plot Theta0(k,tau)
#
ax_Theta = fig.add_subplot(111)
print '> Plotting Theta_0'
# fig_Theta = ax_Theta.pcolormesh(K,T,cs2_fld,cmap='coolwarm',vmin=cs2_fld.min(), vmax=cs2_fld.max()) #,shading='gouraud')
fig_Theta = ax_Theta.pcolormesh(K,T,cs2_fld,cmap='coolwarm',vmin=cs2_fld.min(), vmax=cs2_fld.max()) #,shading='gouraud')
print '> Done'
#
# plot lines (characteristic times and scales)
#
# ax_Theta.axhline(y=tau_rec,color='k',linestyle='-')
# ax_Theta.axhline(y=tau_eq,color='k',linestyle='-')
# ax_Theta.axhline(y=tau_lambda,color='k',linestyle='-')
ax_Theta.axhline(y=tau_ac,color='white',linestyle='-',linewidth=2,zorder=1)
# ax_Theta.axvline(x=k_ac,color='k',linestyle='-')
ax_Theta.plot(aH,tau,'--',color='white',linewidth=2,zorder=1)
# ax_Theta.plot(ks,tau,color='#FFFF33',linestyle='-',linewidth=2)
# ax_Theta.plot(kd,tau,'b-',linewidth=2)
#
# dealing with labels
#
k_range = np.linspace(k_ac,k_of_cs2_at_ac(0.9),100)
print "k_range",k_range
ax_Theta.set_title(r'$n=%d,\Theta_i=%.1f,{\rm log}_{10}(z_c)=%.1f$'%(n_axion,Theta_initial_fld,-1*log10ac), fontsize = 22)
# ax_Theta.text(1.5*k[0],0.9*tau_rec,r'$\mathrm{rec.}$', fontsize = 22)
ax_Theta.text(0.6,0.9*tau_ac,r'$\tau(z_c)$', fontsize = 22,color='white')
# ax_Theta.text(1.5*k[0],0.9*tau_eq,r'$\mathrm{R/M} \,\, \mathrm{eq.}$', fontsize = 22)
# ax_Theta.text(1.5*k[0],0.9*tau_lambda,r'$\mathrm{M/L} \,\, \mathrm{eq.}$', fontsize = 22)
# ax_Theta.text(1.1*k_ac,1.1*tau[0],r'$k=a_cH(a_c)$', fontsize = 22)
# ax_Theta.annotate(r'$\mathrm{Hubble} \,\, \mathrm{cross.}$',
# xy=(background_aH_at_tau(tau_label_Hubble),tau_label_Hubble),
# xytext=(0.1*background_aH_at_tau(tau_label_Hubble),0.8*tau_label_Hubble),
# arrowprops=dict(color='white', shrink=0.05, width=1, headlength=2, headwidth=2), fontsize = 22,color='white')
# ax_Theta.text(0.3,17,r'$\mathrm{Hubble} \,\, \mathrm{cross.}$',fontsize = 22,color='white',rotation=35)
# ax_Theta.text(0.08,17,r'$\mathrm{Hubble}$',fontsize = 22,color='white',rotation=35)
ax_Theta.text(0.04,17,r'$\mathrm{Hubble}$',fontsize = 22,color='white',rotation=35)
# ax_Theta.text(0.3,17,r'$\mathrm{Hubble} \,\, \mathrm{cross.}$',fontsize = 22,color='white',rotation=35)
# xy=(background_aH_at_tau(tau_label_Hubble),tau_label_Hubble),
# xytext=(0.1*background_aH_at_tau(tau_label_Hubble),0.8*tau_label_Hubble),
# arrowprops=dict(color='white', shrink=0.05, width=1, headlength=2, headwidth=2), fontsize = 22,color='white')
# ax_Theta.annotate(r'$\mathrm{sound} \,\, \mathrm{horizon} \,\, \mathrm{cross.}$',
# xy=(background_ks_at_tau(tau_label_ks),tau_label_ks),
# xytext=(0.07*background_aH_at_tau(tau_label_ks),0.8*tau_label_ks),
# arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))
# ax_Theta.annotate(r'$\mathrm{damping} \,\, \mathrm{scale} \,\, \mathrm{cross.}$',
# xy=(thermodynamics_kd_at_tau(tau_label_kd),tau_label_kd),
# xytext=(0.2*thermodynamics_kd_at_tau(tau_label_kd),2.0*tau_label_kd),
# arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))
#
# dealing with axes
#
ax_Theta.set_xlim(k[0],k[-1])
ax_Theta.set_xscale('log')
ax_Theta.set_yscale('log')
ax_Theta.set_xlabel(r'$k \,\,\, \mathrm{[h/Mpc]}$', fontsize = 22)
ax_Theta.set_ylabel(r'$\tau \,\,\, \mathrm{[Mpc]}$', fontsize = 22)
ax_Theta.invert_yaxis()
#
# color legend
#
cbar = fig.colorbar(fig_Theta)
cbar.set_label(r'$c_s^2$', fontsize = 22)
ax_Theta.tick_params(axis='x',labelsize=23)
ax_Theta.tick_params(axis='y',labelsize=23)
ticklabs = cbar.ax.get_yticklabels()
cbar.ax.set_yticklabels(ticklabs, fontsize=20)
contour = ax_Theta.contour(K, T, cs2_fld, Contours_at,
extent = (min(K[0]), max(K[0]), min(T[0]), max(T[0]) ), # place colour map at correct a_c and Om_fld values
origin = 'lower',
corner_mask = True,
colors=('w','w','w','w'),linewidths=('2','2','2','2'))
# manual_locations = [(0.03, 300),(0.03, 40)]
# manual_locations = [(0.03, 300),(0.01, 4000)]
# #2p8:
manual_locations = [(0.006, 200),(0.01, 400),(0.01, 1000)]
#0p1:
# manual_locations = [(0.01, 400),(0.01, 30),(0.05, 50)]
ax_Theta.clabel(contour, fmt='%.2f', colors='w', fontsize=22,ls='--',lw='2',manual=manual_locations,zorder=1)
# ax_Theta.clabel(contour, fmt='%.1f', colors='w', fontsize=22,ls='--',lw='2')
ax_Theta.fill_between(k_range, tau_ac-0.11*tau_ac, tau_ac+0.15*tau_ac, facecolor='blue', alpha=0.5,zorder=2)
plt.savefig('cs2_n3_2p8_vAxiCLASSv2.png',dpi=300)
# plt.savefig('cs2_n3_2p8.png',dpi=300)
# plt.savefig('cs2_n3_2p8.pdf')
else:
fig = plt.figure(figsize=(8,6))
ax_Theta = fig.add_subplot(111)
ax_Theta.plot(Theta_initial_fld_array/np.pi,k9_over_kac,'r',lw=2,label=r'$n=3,~{\rm Log}_{10}(z_c)=3.5$')
ax_Theta.set_xlabel(r'$\Theta_i/\pi$', fontsize = 22)
ax_Theta.set_ylabel(r'$k_{c_s^2<0.9}(z_c)/k_H(z_c)$', fontsize = 22)
ax_Theta.tick_params(axis='x',labelsize=20)
ax_Theta.tick_params(axis='y',labelsize=20)
ax_Theta.legend(frameon=False,prop={'size':18},loc='upper left',borderaxespad=0.)
plt.tight_layout()
plt.savefig('k9_vs_Thetai.png',dpi=300)