mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-05-13 19:41:12 +00:00
update Growth.py to allow using FastPM solver
This commit is contained in:
Wassim Kabalan
parent
e1daa8cba4
commit
20ace41d32
1 changed files with 87 additions and 84 deletions
|
|
@ -119,7 +119,7 @@ def growth_factor(cosmo, a):
|
|||
if cosmo._flags["gamma_growth"]:
|
||||
return _growth_factor_gamma(cosmo, a)
|
||||
else:
|
||||
return _growth_factor_ODE(cosmo, a)
|
||||
return _growth_factor_ODE(cosmo, a)[0]
|
||||
|
||||
|
||||
def growth_factor_second(cosmo, a):
|
||||
|
|
@ -225,7 +225,7 @@ def growth_rate_second(cosmo, a):
|
|||
return _growth_rate_second_ODE(cosmo, a)
|
||||
|
||||
|
||||
def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
|
||||
def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=256, eps=1e-4):
|
||||
"""Compute linear growth factor D(a) at a given scale factor,
|
||||
normalised such that D(a=1) = 1.
|
||||
|
||||
|
|
@ -243,7 +243,7 @@ def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
|
|||
Growth factor computed at requested scale factor
|
||||
"""
|
||||
# Check if growth has already been computed
|
||||
if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
#if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
# Compute tabulated array
|
||||
atab = np.logspace(log10_amin, 0.0, steps)
|
||||
|
||||
|
|
@ -281,6 +281,7 @@ def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
|
|||
htab = dyda2[:, 1, 0] / y1[-1] * atab / gtab
|
||||
h2tab = dyda2[:, 1, 1] / y2[-1] * atab / g2tab
|
||||
|
||||
|
||||
cache = {
|
||||
"a": atab,
|
||||
"g": gtab,
|
||||
|
|
@ -290,10 +291,8 @@ def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
|
|||
"f2": f2tab,
|
||||
"h2": h2tab,
|
||||
}
|
||||
cosmo._workspace["background.growth_factor"] = cache
|
||||
else:
|
||||
cache = cosmo._workspace["background.growth_factor"]
|
||||
return np.clip(interp(a, cache["a"], cache["g"]), 0.0, 1.0)
|
||||
|
||||
return np.clip(interp(a, cache["a"], cache["g"]), 0.0, 1.0) , cache
|
||||
|
||||
|
||||
def _growth_rate_ODE(cosmo, a):
|
||||
|
|
@ -314,11 +313,9 @@ def _growth_rate_ODE(cosmo, a):
|
|||
Growth rate computed at requested scale factor
|
||||
"""
|
||||
# Check if growth has already been computed, if not, compute it
|
||||
if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
_growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = cosmo._workspace["background.growth_factor"]
|
||||
return interp(a, cache["a"], cache["f"])
|
||||
|
||||
cache = _growth_factor_ODE(cosmo, np.atleast_1d(1.0))[1]
|
||||
return interp(a, cache["a"], cache["f"])
|
||||
|
||||
def _growth_factor_second_ODE(cosmo, a):
|
||||
"""Compute second order growth factor D2(a) at a given scale factor,
|
||||
|
|
@ -338,36 +335,12 @@ def _growth_factor_second_ODE(cosmo, a):
|
|||
Second order growth factor computed at requested scale factor
|
||||
"""
|
||||
# Check if growth has already been computed, if not, compute it
|
||||
if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
_growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = cosmo._workspace["background.growth_factor"]
|
||||
#if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
# _growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = _growth_factor_ODE(cosmo, a)[1]
|
||||
return interp(a, cache["a"], cache["g2"])
|
||||
|
||||
|
||||
def _growth_rate_ODE(cosmo, a):
|
||||
"""Compute growth rate dD/dlna at a given scale factor by solving the linear
|
||||
growth ODE.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cosmo: `Cosmology`
|
||||
Cosmology object
|
||||
|
||||
a: array_like
|
||||
Scale factor
|
||||
|
||||
Returns
|
||||
-------
|
||||
f: ndarray, or float if input scalar
|
||||
Second order growth rate computed at requested scale factor
|
||||
"""
|
||||
# Check if growth has already been computed, if not, compute it
|
||||
if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
_growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = cosmo._workspace["background.growth_factor"]
|
||||
return interp(a, cache["a"], cache["f"])
|
||||
|
||||
|
||||
def _growth_rate_second_ODE(cosmo, a):
|
||||
"""Compute second order growth rate dD2/dlna at a given scale factor by solving the linear
|
||||
growth ODE.
|
||||
|
|
@ -386,9 +359,9 @@ def _growth_rate_second_ODE(cosmo, a):
|
|||
Second order growth rate computed at requested scale factor
|
||||
"""
|
||||
# Check if growth has already been computed, if not, compute it
|
||||
if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
_growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = cosmo._workspace["background.growth_factor"]
|
||||
#if not "background.growth_factor" in cosmo._workspace.keys():
|
||||
# _growth_factor_ODE(cosmo, np.atleast_1d(1.0))
|
||||
cache = _growth_factor_ODE(cosmo, a)[1]
|
||||
return interp(a, cache["a"], cache["f2"])
|
||||
|
||||
|
||||
|
|
@ -521,6 +494,34 @@ def Gf2(cosmo, a):
|
|||
return D2f * np.power(a, 3) * np.power(Esqr(cosmo, a), 0.5)
|
||||
|
||||
|
||||
def gp(cosmo, a):
|
||||
r""" Derivative of D1 against a
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cosmo: dict
|
||||
Cosmology dictionary.
|
||||
|
||||
a : array_like
|
||||
Scale factor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Scalar float Tensor : the derivative of D1 against a.
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
The expression for :math:`gp(a)` is:
|
||||
|
||||
.. math::
|
||||
gp(a)=\frac{dD1}{da}= D'_{1norm}/a
|
||||
"""
|
||||
f1 = growth_rate(cosmo, a)
|
||||
g1 = growth_factor(cosmo, a)
|
||||
D1f = f1 * g1 / a
|
||||
return D1f
|
||||
|
||||
def dGfa(cosmo, a):
|
||||
r""" Derivative of Gf against a
|
||||
|
||||
|
|
@ -549,7 +550,8 @@ def dGfa(cosmo, a):
|
|||
f1 = growth_rate(cosmo, a)
|
||||
g1 = growth_factor(cosmo, a)
|
||||
D1f = f1 * g1 / a
|
||||
cache = cosmo._workspace['background.growth_factor']
|
||||
#cache = cosmo._workspace['background.growth_factor']
|
||||
cache = _growth_factor_ODE(cosmo, a)[1]
|
||||
f1p = cache['h'] / cache['a'] * cache['g']
|
||||
f1p = interp(np.log(a), np.log(cache['a']), f1p)
|
||||
Ea = E(cosmo, a)
|
||||
|
|
@ -584,7 +586,8 @@ def dGf2a(cosmo, a):
|
|||
f2 = growth_rate_second(cosmo, a)
|
||||
g2 = growth_factor_second(cosmo, a)
|
||||
D2f = f2 * g2 / a
|
||||
cache = cosmo._workspace['background.growth_factor']
|
||||
#cache = cosmo._workspace['background.growth_factor']
|
||||
cache = _growth_factor_ODE(cosmo, a)[1]
|
||||
f2p = cache['h2'] / cache['a'] * cache['g2']
|
||||
f2p = interp(np.log(a), np.log(cache['a']), f2p)
|
||||
E_a = E(cosmo, a)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue