update Growth.py to allow using FastPM solver

jaxpm/growth.py

@@ -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:
+    return np.clip(interp(a, cache["a"], cache["g"]), 0.0, 1.0) , cache
-        cache = cosmo._workspace["background.growth_factor"]
-    return np.clip(interp(a, cache["a"], cache["g"]), 0.0, 1.0)
 
 
 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():
+    #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"])
 
 
-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():
+    #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"])
 
 
@@ -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)