mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-05-14 03:51:11 +00:00
Applying formatting
This commit is contained in:
EiffL
parent
835fa89aec
commit
f28442bb48
14 changed files with 565 additions and 445 deletions
|
|
@ -1,8 +1,8 @@
|
|||
import jax.numpy as np
|
||||
|
||||
from jax_cosmo.background import *
|
||||
from jax_cosmo.scipy.interpolate import interp
|
||||
from jax_cosmo.scipy.ode import odeint
|
||||
from jax_cosmo.background import *
|
||||
|
||||
|
||||
def E(cosmo, a):
|
||||
r"""Scale factor dependent factor E(a) in the Hubble
|
||||
|
|
@ -52,12 +52,8 @@ def df_de(cosmo, a, epsilon=1e-5):
|
|||
\frac{df}{da}(a) = =\frac{3w_a \left( \ln(a-\epsilon)-
|
||||
\frac{a-1}{a-\epsilon}\right)}{\ln^2(a-\epsilon)}
|
||||
"""
|
||||
return (
|
||||
3
|
||||
* cosmo.wa
|
||||
* (np.log(a - epsilon) - (a - 1) / (a - epsilon))
|
||||
/ np.power(np.log(a - epsilon), 2)
|
||||
)
|
||||
return (3 * cosmo.wa * (np.log(a - epsilon) - (a - 1) / (a - epsilon)) /
|
||||
np.power(np.log(a - epsilon), 2))
|
||||
|
||||
|
||||
def dEa(cosmo, a):
|
||||
|
|
@ -89,15 +85,11 @@ def dEa(cosmo, a):
|
|||
where :math:`f(a)` is the Dark Energy evolution parameter computed
|
||||
by :py:meth:`.f_de`.
|
||||
"""
|
||||
return (
|
||||
0.5
|
||||
* (
|
||||
-3 * cosmo.Omega_m * np.power(a, -4)
|
||||
- 2 * cosmo.Omega_k * np.power(a, -3)
|
||||
+ df_de(cosmo, a) * cosmo.Omega_de * np.power(a, f_de(cosmo, a))
|
||||
)
|
||||
/ np.power(Esqr(cosmo, a), 0.5)
|
||||
)
|
||||
return (0.5 *
|
||||
(-3 * cosmo.Omega_m * np.power(a, -4) -
|
||||
2 * cosmo.Omega_k * np.power(a, -3) +
|
||||
df_de(cosmo, a) * cosmo.Omega_de * np.power(a, f_de(cosmo, a))) /
|
||||
np.power(Esqr(cosmo, a), 0.5))
|
||||
|
||||
|
||||
def growth_factor(cosmo, a):
|
||||
|
|
@ -155,8 +147,7 @@ def growth_factor_second(cosmo, a):
|
|||
"""
|
||||
if cosmo._flags["gamma_growth"]:
|
||||
raise NotImplementedError(
|
||||
"Gamma growth rate is not implemented for second order growth!"
|
||||
)
|
||||
"Gamma growth rate is not implemented for second order growth!")
|
||||
return None
|
||||
else:
|
||||
return _growth_factor_second_ODE(cosmo, a)
|
||||
|
|
@ -228,8 +219,7 @@ def growth_rate_second(cosmo, a):
|
|||
"""
|
||||
if cosmo._flags["gamma_growth"]:
|
||||
raise NotImplementedError(
|
||||
"Gamma growth factor is not implemented for second order growth!"
|
||||
)
|
||||
"Gamma growth factor is not implemented for second order growth!")
|
||||
return None
|
||||
else:
|
||||
return _growth_rate_second_ODE(cosmo, a)
|
||||
|
|
@ -258,23 +248,19 @@ def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
|
|||
atab = np.logspace(log10_amin, 0.0, steps)
|
||||
|
||||
def D_derivs(y, x):
|
||||
q = (
|
||||
2.0
|
||||
- 0.5
|
||||
* (
|
||||
Omega_m_a(cosmo, x)
|
||||
+ (1.0 + 3.0 * w(cosmo, x)) * Omega_de_a(cosmo, x)
|
||||
)
|
||||
) / x
|
||||
q = (2.0 - 0.5 *
|
||||
(Omega_m_a(cosmo, x) +
|
||||
(1.0 + 3.0 * w(cosmo, x)) * Omega_de_a(cosmo, x))) / x
|
||||
r = 1.5 * Omega_m_a(cosmo, x) / x / x
|
||||
|
||||
g1, g2 = y[0]
|
||||
f1, f2 = y[1]
|
||||
dy1da = [f1, -q * f1 + r * g1]
|
||||
dy2da = [f2, -q * f2 + r * g2 - r * g1 ** 2]
|
||||
dy2da = [f2, -q * f2 + r * g2 - r * g1**2]
|
||||
return np.array([[dy1da[0], dy2da[0]], [dy1da[1], dy2da[1]]])
|
||||
|
||||
y0 = np.array([[atab[0], -3.0 / 7 * atab[0] ** 2], [1.0, -6.0 / 7 * atab[0]]])
|
||||
y0 = np.array([[atab[0], -3.0 / 7 * atab[0]**2],
|
||||
[1.0, -6.0 / 7 * atab[0]]])
|
||||
y = odeint(D_derivs, y0, atab)
|
||||
|
||||
# compute second order derivatives growth
|
||||
|
|
@ -473,8 +459,7 @@ def _growth_rate_gamma(cosmo, a):
|
|||
|
||||
see :cite:`2019:Euclid Preparation VII, eqn.32`
|
||||
"""
|
||||
return Omega_m_a(cosmo, a) ** cosmo.gamma
|
||||
|
||||
return Omega_m_a(cosmo, a)**cosmo.gamma
|
||||
|
||||
|
||||
def Gf(cosmo, a):
|
||||
|
|
@ -503,7 +488,7 @@ def Gf(cosmo, a):
|
|||
"""
|
||||
f1 = growth_rate(cosmo, a)
|
||||
g1 = growth_factor(cosmo, a)
|
||||
D1f = f1*g1/ a
|
||||
D1f = f1 * g1 / a
|
||||
return D1f * np.power(a, 3) * np.power(Esqr(cosmo, a), 0.5)
|
||||
|
||||
|
||||
|
|
@ -532,7 +517,7 @@ def Gf2(cosmo, a):
|
|||
"""
|
||||
f2 = growth_rate_second(cosmo, a)
|
||||
g2 = growth_factor_second(cosmo, a)
|
||||
D2f = f2*g2/ a
|
||||
D2f = f2 * g2 / a
|
||||
return D2f * np.power(a, 3) * np.power(Esqr(cosmo, a), 0.5)
|
||||
|
||||
|
||||
|
|
@ -563,13 +548,12 @@ def dGfa(cosmo, a):
|
|||
"""
|
||||
f1 = growth_rate(cosmo, a)
|
||||
g1 = growth_factor(cosmo, a)
|
||||
D1f = f1*g1/ a
|
||||
D1f = f1 * g1 / a
|
||||
cache = cosmo._workspace['background.growth_factor']
|
||||
f1p = cache['h'] / cache['a'] * cache['g']
|
||||
f1p = interp(np.log(a), np.log(cache['a']), f1p)
|
||||
Ea = E(cosmo, a)
|
||||
return (f1p * a**3 * Ea + D1f * a**3 * dEa(cosmo, a) +
|
||||
3 * a**2 * Ea * D1f)
|
||||
return (f1p * a**3 * Ea + D1f * a**3 * dEa(cosmo, a) + 3 * a**2 * Ea * D1f)
|
||||
|
||||
|
||||
def dGf2a(cosmo, a):
|
||||
|
|
@ -599,10 +583,9 @@ def dGf2a(cosmo, a):
|
|||
"""
|
||||
f2 = growth_rate_second(cosmo, a)
|
||||
g2 = growth_factor_second(cosmo, a)
|
||||
D2f = f2*g2/ a
|
||||
D2f = f2 * g2 / a
|
||||
cache = cosmo._workspace['background.growth_factor']
|
||||
f2p = cache['h2'] / cache['a'] * cache['g2']
|
||||
f2p = interp(np.log(a), np.log(cache['a']), f2p)
|
||||
E = E(cosmo, a)
|
||||
return (f2p * a**3 * E + D2f * a**3 * dEa(cosmo, a) +
|
||||
3 * a**2 * E * D2f)
|
||||
return (f2p * a**3 * E + D2f * a**3 * dEa(cosmo, a) + 3 * a**2 * E * D2f)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue