guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-05-13 19:41:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

update Growth.py to allow using FastPM solver

Browse source
This commit is contained in:
Wassim Kabalan 2025-05-08 15:08:28 +02:00
parent e1daa8cba4
commit 20ace41d32
1 changed files with 87 additions and 84 deletions
Download patch file Download diff file Expand all files Collapse all files

171
jaxpm/growth.py
View file

@ -119,7 +119,7 @@ def growth_factor(cosmo, a):
if cosmo._flags["gamma_growth"]: if cosmo._flags["gamma_growth"]:
return _growth_factor_gamma(cosmo, a) return _growth_factor_gamma(cosmo, a)
else: else:
return _growth_factor_ODE(cosmo, a) return _growth_factor_ODE(cosmo, a)[0]
def growth_factor_second(cosmo, a): def growth_factor_second(cosmo, a):
@ -225,7 +225,7 @@ def growth_rate_second(cosmo, a):
return _growth_rate_second_ODE(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, """Compute linear growth factor D(a) at a given scale factor,
normalised such that D(a=1) = 1. normalised such that D(a=1) = 1.
@ -243,57 +243,56 @@ def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
Growth factor computed at requested scale factor Growth factor computed at requested scale factor
""" """
# Check if growth has already been computed # 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 # Compute tabulated array
atab = np.logspace(log10_amin, 0.0, steps) atab = np.logspace(log10_amin, 0.0, steps)
def D_derivs(y, x): def D_derivs(y, x):
q = (2.0 - 0.5 * q = (2.0 - 0.5 *
(Omega_m_a(cosmo, x) + (Omega_m_a(cosmo, x) +
(1.0 + 3.0 * w(cosmo, x)) * Omega_de_a(cosmo, x))) / x (1.0 + 3.0 * w(cosmo, x)) * Omega_de_a(cosmo, x))) / x
r = 1.5 * Omega_m_a(cosmo, x) / x / x r = 1.5 * Omega_m_a(cosmo, x) / x / x
g1, g2 = y[0] g1, g2 = y[0]
f1, f2 = y[1] f1, f2 = y[1]
dy1da = [f1, -q * f1 + r * g1] 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]]]) return np.array([[dy1da[0], dy2da[0]], [dy1da[1], dy2da[1]]])
y0 = np.array([[atab[0], -3.0 / 7 * atab[0]**2], y0 = np.array([[atab[0], -3.0 / 7 * atab[0]**2],
[1.0, -6.0 / 7 * atab[0]]]) [1.0, -6.0 / 7 * atab[0]]])
y = odeint(D_derivs, y0, atab) y = odeint(D_derivs, y0, atab)
# compute second order derivatives growth # compute second order derivatives growth
dyda2 = D_derivs(np.transpose(y, (1, 2, 0)), atab) dyda2 = D_derivs(np.transpose(y, (1, 2, 0)), atab)
dyda2 = np.transpose(dyda2, (2, 0, 1)) dyda2 = np.transpose(dyda2, (2, 0, 1))
# Normalize results # Normalize results
y1 = y[:, 0, 0] y1 = y[:, 0, 0]
gtab = y1 / y1[-1] gtab = y1 / y1[-1]
y2 = y[:, 0, 1] y2 = y[:, 0, 1]
g2tab = y2 / y2[-1] g2tab = y2 / y2[-1]
# To transform from dD/da to dlnD/dlna: dlnD/dlna = a / D dD/da # To transform from dD/da to dlnD/dlna: dlnD/dlna = a / D dD/da
ftab = y[:, 1, 0] / y1[-1] * atab / gtab ftab = y[:, 1, 0] / y1[-1] * atab / gtab
f2tab = y[:, 1, 1] / y2[-1] * atab / g2tab f2tab = y[:, 1, 1] / y2[-1] * atab / g2tab
# Similarly for second order derivatives # Similarly for second order derivatives
# Note: these factors are not accessible as parent functions yet # Note: these factors are not accessible as parent functions yet
# since it is unclear what to refer to them with. # since it is unclear what to refer to them with.
htab = dyda2[:, 1, 0] / y1[-1] * atab / gtab htab = dyda2[:, 1, 0] / y1[-1] * atab / gtab
h2tab = dyda2[:, 1, 1] / y2[-1] * atab / g2tab h2tab = dyda2[:, 1, 1] / y2[-1] * atab / g2tab
cache = {
"a": atab, cache = {
"g": gtab, "a": atab,
"f": ftab, "g": gtab,
"h": htab, "f": ftab,
"g2": g2tab, "h": htab,
"f2": f2tab, "g2": g2tab,
"h2": h2tab, "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) , cache
return np.clip(interp(a, cache["a"], cache["g"]), 0.0, 1.0)
def _growth_rate_ODE(cosmo, a): def _growth_rate_ODE(cosmo, a):
@ -314,12 +313,10 @@ def _growth_rate_ODE(cosmo, a):
Growth rate computed at requested scale factor Growth rate computed at requested scale factor
""" """
# Check if growth has already been computed, if not, compute it # 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 = _growth_factor_ODE(cosmo, np.atleast_1d(1.0))[1]
cache = cosmo._workspace["background.growth_factor"]
return interp(a, cache["a"], cache["f"]) return interp(a, cache["a"], cache["f"])
def _growth_factor_second_ODE(cosmo, a): def _growth_factor_second_ODE(cosmo, a):
"""Compute second order growth factor D2(a) at a given scale factor, """Compute second order growth factor D2(a) at a given scale factor,
normalised such that D(a=1) = 1. normalised such that D(a=1) = 1.
@ -338,36 +335,12 @@ def _growth_factor_second_ODE(cosmo, a):
Second order growth factor computed at requested scale factor Second order growth factor computed at requested scale factor
""" """
# Check if growth has already been computed, if not, compute it # Check if growth has already been computed, if not, compute it
if not "background.growth_factor" in cosmo._workspace.keys(): #if not "background.growth_factor" in cosmo._workspace.keys():
_growth_factor_ODE(cosmo, np.atleast_1d(1.0)) # _growth_factor_ODE(cosmo, np.atleast_1d(1.0))
cache = cosmo._workspace["background.growth_factor"] cache = _growth_factor_ODE(cosmo, a)[1]
return interp(a, cache["a"], cache["g2"]) 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): def _growth_rate_second_ODE(cosmo, a):
"""Compute second order growth rate dD2/dlna at a given scale factor by solving the linear """Compute second order growth rate dD2/dlna at a given scale factor by solving the linear
growth ODE. growth ODE.
@ -386,9 +359,9 @@ def _growth_rate_second_ODE(cosmo, a):
Second order growth rate computed at requested scale factor Second order growth rate computed at requested scale factor
""" """
# Check if growth has already been computed, if not, compute it # Check if growth has already been computed, if not, compute it
if not "background.growth_factor" in cosmo._workspace.keys(): #if not "background.growth_factor" in cosmo._workspace.keys():
_growth_factor_ODE(cosmo, np.atleast_1d(1.0)) # _growth_factor_ODE(cosmo, np.atleast_1d(1.0))
cache = cosmo._workspace["background.growth_factor"] cache = _growth_factor_ODE(cosmo, a)[1]
return interp(a, cache["a"], cache["f2"]) 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) 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): def dGfa(cosmo, a):
r""" Derivative of Gf against a r""" Derivative of Gf against a
@ -549,7 +550,8 @@ def dGfa(cosmo, a):
f1 = growth_rate(cosmo, a) f1 = growth_rate(cosmo, a)
g1 = growth_factor(cosmo, a) g1 = growth_factor(cosmo, a)
D1f = f1 * g1 / 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 = cache['h'] / cache['a'] * cache['g']
f1p = interp(np.log(a), np.log(cache['a']), f1p) f1p = interp(np.log(a), np.log(cache['a']), f1p)
Ea = E(cosmo, a) Ea = E(cosmo, a)
@ -584,9 +586,10 @@ def dGf2a(cosmo, a):
f2 = growth_rate_second(cosmo, a) f2 = growth_rate_second(cosmo, a)
g2 = growth_factor_second(cosmo, a) g2 = growth_factor_second(cosmo, a)
D2f = f2 * g2 / 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 = cache['h2'] / cache['a'] * cache['g2']
f2p = interp(np.log(a), np.log(cache['a']), f2p) f2p = interp(np.log(a), np.log(cache['a']), f2p)
E_a = E(cosmo, a) E_a = E(cosmo, a)
return (f2p * a**3 * E_a + D2f * a**3 * dEa(cosmo, a) + return (f2p * a**3 * E_a + D2f * a**3 * dEa(cosmo, a) +
3 * a**2 * E_a * D2f) 3 * a**2 * E_a * D2f)