-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathgen_cosmo_fcns.py
166 lines (144 loc) · 5.57 KB
/
gen_cosmo_fcns.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
from __future__ import print_function, division
import numpy as np
from scipy import interpolate
from collections import OrderedDict
def generate_calc_Da(test_plot=False,
N_integration=10000,
cosmo=None,
verbose=True):
"""
Return the function calc_Da, which takes a as argument
and returns D(a).
"""
if verbose:
print("Compute D(a)")
a = np.linspace(1.0e-4, 1.0, N_integration)
H_over_H0 = np.sqrt(cosmo.Om_r / a**4 + cosmo.Om_m / a**3 + cosmo.Om_L +
cosmo.Om_K / a**2)
Da = np.zeros(a.shape)
# compute the integral
for imax, aa in enumerate(a):
Da[imax] = np.trapz(1.0 / (a[:imax + 1] * H_over_H0[:imax + 1])**3,
a[:imax + 1])
# Prefactors
Da = Da * 5. / 2. * cosmo.Om_m * H_over_H0
if verbose:
print("Got D(a)")
def calc_Da(aeval):
return np.interp(aeval, a, Da)
# Plot
if test_plot:
import matplotlib.pyplot as plt
print("Making growth.pdf")
fig, ax = plt.subplots(1, 1)
ax.plot(a, calc_Da(a), 'b-')
ax.plot(a, a, 'b:')
ax.set_xlabel('$a$')
ax.set_ylabel('$D(a)$')
plt.tight_layout()
plt.savefig("growth.pdf")
#plt.show()
return calc_Da
def calc_f_log_growth_rate(a=None, calc_Da=None, cosmo=None, do_test=False):
"""Calculate f=dlnD/dlna"""
# derived from formula for D
assert cosmo.Om_K == 0.
H_over_H0 = np.sqrt(cosmo.Om_r / a**4 + cosmo.Om_m / a**3 + cosmo.Om_L +
cosmo.Om_K / a**2)
f_analytical = 1. / (a * H_over_H0)**2 * (-2. * cosmo.Om_r / a**2 -
3. / 2. * cosmo.Om_m / a +
5. / 2. * cosmo.Om_m / calc_Da(a))
if do_test:
# compute numerically and compare
avec = np.linspace(1.0e-4, 1.0, 1000)
d_D_d_a = np.interp(a, 0.5 * (avec[1:] + avec[:-1]),
np.diff(calc_Da(avec)) / np.diff(avec))
f_numerical = a / calc_Da(a) * d_D_d_a
print('f_numerical=%g' % f_numerical)
print('f_numerical/f_analytical-1 = %g' %
(f_numerical / f_analytical - 1.))
assert np.isclose(f_numerical, f_analytical, rtol=1e-4)
return f_analytical
def generate_calc_chi(test_plot=False, N_integration=10000, cosmo=None):
"""
Return the function calc_chiMpc_a, which takes a as argument
and returns chi(a) in Mpc.
"""
print("Compute chi(a)")
# This must be linear in a (assume this for gradient below!)
a = np.linspace(1.0e-4, 1.0, N_integration)
H_over_H0 = np.sqrt(cosmo.Om_r / a**4 + cosmo.Om_m / a**3 + cosmo.Om_L +
cosmo.Om_K / a**2)
chi_a = np.zeros(a.shape)
# compute the integral to get H0*chi(a) = int_a^1 da/(a^2 H(a)/H0)
for imin, aa in enumerate(a):
chi_a[imin] = np.trapz(1.0 / (a[imin:]**2 * H_over_H0[imin:]), a[imin:])
# Prefactors: chi = [H0*chi]/H0
chi_a /= (100. * cosmo.h0
) # this is chi in units of 1/(km/s/Mpc)=Mpc/(km/s)
# Multiply by c in km/s to get chi in Mpc
chi_a *= 2.99792458e5
print("Got chi(a)")
def calc_chiMpc_a(aeval):
"""
Calculate chi(a) in Mpc.
"""
#if type(aeval)==list:
# aeval = np.array(aeval)
return np.interp(aeval, a, chi_a)
# Also compute derivative dchi/da
# TODO: could also get this easier from 1/(a**2 H(a))
grad_a = np.gradient(chi_a, a[1] - a[0])
def calc_delchiMpc_dela_a(aeval):
return np.interp(aeval, a, grad_a)
from scipy.optimize import root
def calc_a_chiMpc(chiMpc_eval):
"""
Invert chi(a) to get a(chi). Assumes chi in Mpc.
"""
if type(chiMpc_eval) == np.ndarray:
aout = np.zeros(chiMpc_eval.shape) + np.nan
#print("Chi eval:", chiMpc_eval)
for cc, this_chiMpc in enumerate(chiMpc_eval):
#print("this chi", cc, this_chiMpc, calc_a_chiMpc(this_chiMpc))
aout[cc] = calc_a_chiMpc(this_chiMpc)
return aout
else:
def rootfcn(a):
return calc_chiMpc_a(a) - chiMpc_eval
soln = root(rootfcn, 0.5)
return soln.x[0]
# Plot
if test_plot:
import matplotlib.pyplot as plt
print("Making chi.pdf")
fig, axlst = plt.subplots(1, 2, figsize=(16, 8), sharey=True)
for i, zmax in enumerate([20, 1100]):
z = 1. / a - 1
ax = axlst[i]
ax.plot(z, calc_chiMpc_a(a), 'b-', label="$\chi$")
# matter domination in flat universe (Dodelson Eq. 2.43)
ax.plot(z,
2. / (100. * cosmo.h0 / 2.9979e5) * (1. - a**0.5),
'b:',
label="$\\frac{2}{H_0} (1-a^{1/2})$, matter dom.")
ax.plot(z,
0. * z + 2. / (100. * cosmo.h0 / 2.9979e5),
'r:',
label="$\\frac{2}{H_0}$")
ax.set_xlabel('$z$')
ax.set_xlim((0, zmax))
axlst[0].set_ylabel('$\chi(z)$ [Mpc]')
axlst[-1].legend(loc="best")
plt.tight_layout()
plt.savefig("chi_a.pdf")
#plt.show()
# plot d chi/ d a
fig, ax = plt.subplots(1, 1)
atmp = np.linspace(1.0e-4, 1.0, 333)
ax.plot(atmp, calc_delchiMpc_dela_a(atmp))
ax.set_ylabel('$d\chi/d a$')
ax.set_xlabel('$a$')
plt.tight_layout()
plt.savefig("dchi_da.pdf")
return calc_chiMpc_a, calc_delchiMpc_dela_a, calc_a_chiMpc