JaxPM/jaxpm/growth.py

import jax.numpy as np
from jax.numpy import interp
from jax_cosmo.background import *
from jax_cosmo.scipy.ode import odeint


def E(cosmo, a):
    r"""Scale factor dependent factor E(a) in the Hubble
    parameter.
    Parameters
    ----------
    a : array_like
        Scale factor
    Returns
    -------
    E : ndarray, or float if input scalar
        Square of the scaling of the Hubble constant as a function of
        scale factor
    Notes
    -----
    The Hubble parameter at scale factor `a` is given by
    :math:`H^2(a) = E^2(a) H_o^2` where :math:`E^2` is obtained through
    Friedman's Equation (see :cite:`2005:Percival`) :
    .. math::
        E^2(a) = \Omega_m a^{-3} + \Omega_k a^{-2} + \Omega_{de} a^{f(a)}
    where :math:`f(a)` is the Dark Energy evolution parameter computed
    by :py:meth:`.f_de`.
    """
    return np.power(Esqr(cosmo, a), 0.5)


def df_de(cosmo, a, epsilon=1e-5):
    r"""Derivative of the evolution parameter for the Dark Energy density
    f(a) with respect to the scale factor.
    Parameters
    ----------
    cosmo: Cosmology
    Cosmological parameters structure
    a : array_like
    Scale factor
    epsilon: float value
    Small number to make sure we are not dividing by 0 and avoid a singularity
    Returns
    -------
    df(a)/da : ndarray, or float if input scalar
    Derivative of the evolution parameter for the Dark Energy density
    with respect to the scale factor.
    Notes
    -----
    The expression for :math:`\frac{df(a)}{da}` is:
    .. math::
    \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))


def dEa(cosmo, a):
    r"""Derivative of the scale factor dependent factor E(a) in the Hubble
    parameter.
    Parameters
    ----------
    a : array_like
        Scale factor
    Returns
    -------
    dE(a)/da : ndarray, or float if input scalar
        Derivative of the scale factor dependent factor in the Hubble
      parameter with respect to the scale factor.
    Notes
    -----
    The expression for :math:`\frac{dE}{da}` is:
    .. math::
        \frac{dE(a)}{da}=\frac{-3a^{-4}\Omega_{0m}
        -2a^{-3}\Omega_{0k}
        +f'_{de}\Omega_{0de}a^{f_{de}(a)}}{2E(a)}
    Notes
    -----
    The Hubble parameter at scale factor `a` is given by
    :math:`H^2(a) = E^2(a) H_o^2` where :math:`E^2` is obtained through
    Friedman's Equation (see :cite:`2005:Percival`) :
    .. math::
        E^2(a) = \Omega_m a^{-3} + \Omega_k a^{-2} + \Omega_{de} a^{f(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))


def growth_factor(cosmo, a):
    """Compute linear growth factor D(a) at a given scale factor,
    normalized such that D(a=1) = 1.

    Parameters
    ----------
    cosmo: `Cosmology`
      Cosmology object

    a: array_like
      Scale factor

    Returns
    -------
    D:  ndarray, or float if input scalar
        Growth factor computed at requested scale factor

    Notes
    -----
    The growth computation will depend on the cosmology parametrization, for
    instance if the $\gamma$ parameter is defined, the growth will be computed
    assuming the $f = \Omega^\gamma$ growth rate, otherwise the usual ODE for
    growth will be solved.
    """
    if cosmo._flags["gamma_growth"]:
        return _growth_factor_gamma(cosmo, a)
    else:
        return _growth_factor_ODE(cosmo, a)


def growth_factor_second(cosmo, a):
    """Compute second order growth factor D2(a) at a given scale factor,
    normalized such that D(a=1) = 1.

    Parameters
    ----------
    cosmo: `Cosmology`
      Cosmology object

    a: array_like
      Scale factor

    Returns
    -------
    D2:  ndarray, or float if input scalar
        Growth factor computed at requested scale factor

    Notes
    -----
    The growth computation will depend on the cosmology parametrization,
    as for the linear growth. Currently the second order growth
    factor is not implemented with $\gamma$ parameter.
    """
    if cosmo._flags["gamma_growth"]:
        raise NotImplementedError(
            "Gamma growth rate is not implemented for second order growth!")
        return None
    else:
        return _growth_factor_second_ODE(cosmo, a)


def growth_rate(cosmo, a):
    """Compute growth rate dD/dlna at a given scale factor.

    Parameters
    ----------
    cosmo: `Cosmology`
      Cosmology object

    a: array_like
      Scale factor

    Returns
    -------
    f:  ndarray, or float if input scalar
        Growth rate computed at requested scale factor

    Notes
    -----
    The growth computation will depend on the cosmology parametrization, for
    instance if the $\gamma$ parameter is defined, the growth will be computed
    assuming the $f = \Omega^\gamma$ growth rate, otherwise the usual ODE for
    growth will be solved.

    The LCDM approximation to the growth rate :math:`f_{\gamma}(a)` is given by:

    .. math::

        f_{\gamma}(a) = \Omega_m^{\gamma} (a)

     with :math: `\gamma` in LCDM, given approximately by:
     .. math::

        \gamma = 0.55

    see :cite:`2019:Euclid Preparation VII, eqn.32`
    """
    if cosmo._flags["gamma_growth"]:
        return _growth_rate_gamma(cosmo, a)
    else:
        return _growth_rate_ODE(cosmo, a)


def growth_rate_second(cosmo, a):
    """Compute second order growth rate dD2/dlna at a given scale factor.

    Parameters
    ----------
    cosmo: `Cosmology`
      Cosmology object

    a: array_like
      Scale factor

    Returns
    -------
    f2:  ndarray, or float if input scalar
        Second order growth rate computed at requested scale factor

    Notes
    -----
    The growth computation will depend on the cosmology parametrization,
    as for the linear growth rate. Currently the second order growth
    rate is not implemented with $\gamma$ parameter.
    """
    if cosmo._flags["gamma_growth"]:
        raise NotImplementedError(
            "Gamma growth factor is not implemented for second order growth!")
        return None
    else:
        return _growth_rate_second_ODE(cosmo, a)


def _growth_factor_ODE(cosmo, a, log10_amin=-3, steps=128, eps=1e-4):
    """Compute linear growth factor D(a) at a given scale factor,
    normalised such that D(a=1) = 1.

    Parameters
    ----------
    a: array_like
      Scale factor

    amin: float
      Mininum scale factor, default 1e-3

    Returns
    -------
    D:  ndarray, or float if input scalar
        Growth factor computed at requested scale factor
    """
    # Check if growth has already been computed
    if not "background.growth_factor" in cosmo._workspace.keys():
        # Compute tabulated array
        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
            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]
            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]]])
        y = odeint(D_derivs, y0, atab)

        # compute second order derivatives growth
        dyda2 = D_derivs(np.transpose(y, (1, 2, 0)), atab)
        dyda2 = np.transpose(dyda2, (2, 0, 1))

        # Normalize results
        y1 = y[:, 0, 0]
        gtab = y1 / y1[-1]
        y2 = y[:, 0, 1]
        g2tab = y2 / y2[-1]
        # To transform from dD/da to dlnD/dlna: dlnD/dlna = a / D dD/da
        ftab = y[:, 1, 0] / y1[-1] * atab / gtab
        f2tab = y[:, 1, 1] / y2[-1] * atab / g2tab
        # Similarly for second order derivatives
        # Note: these factors are not accessible as parent functions yet
        # since it is unclear what to refer to them with.
        htab = dyda2[:, 1, 0] / y1[-1] * atab / gtab
        h2tab = dyda2[:, 1, 1] / y2[-1] * atab / g2tab

        cache = {
            "a": atab,
            "g": gtab,
            "f": ftab,
            "h": htab,
            "g2": g2tab,
            "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)


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
        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_factor_second_ODE(cosmo, a):
    """Compute second order growth factor D2(a) at a given scale factor,
    normalised such that D(a=1) = 1.

    Parameters
    ----------
    a: array_like
      Scale factor

    amin: float
      Mininum scale factor, default 1e-3

    Returns
    -------
    D2:  ndarray, or float if input scalar
        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"]
    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.

    Parameters
    ----------
    cosmo: `Cosmology`
      Cosmology object

    a: array_like
      Scale factor

    Returns
    -------
    f2:  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["f2"])


def _growth_factor_gamma(cosmo, a, log10_amin=-3, steps=128):
    r"""Computes growth factor by integrating the growth rate provided by the
    \gamma parametrization. Normalized such that D( a=1) =1

    Parameters
    ----------
    a: array_like
      Scale factor

    amin: float
      Mininum scale factor, default 1e-3

    Returns
    -------
    D:  ndarray, or float if input scalar
        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():
        # Compute tabulated array
        atab = np.logspace(log10_amin, 0.0, steps)

        def integrand(y, loga):
            xa = np.exp(loga)
            return _growth_rate_gamma(cosmo, xa)

        gtab = np.exp(odeint(integrand, np.log(atab[0]), np.log(atab)))
        gtab = gtab / gtab[-1]  # Normalize to a=1.
        cache = {"a": atab, "g": gtab}
        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)


def _growth_rate_gamma(cosmo, a):
    r"""Growth rate approximation at scale factor `a`.

    Parameters
    ----------
    cosmo: `Cosmology`
        Cosmology object

    a : array_like
        Scale factor

    Returns
    -------
    f_gamma : ndarray, or float if input scalar
        Growth rate approximation at the requested scale factor

    Notes
    -----
    The LCDM approximation to the growth rate :math:`f_{\gamma}(a)` is given by:

    .. math::

        f_{\gamma}(a) = \Omega_m^{\gamma} (a)

     with :math: `\gamma` in LCDM, given approximately by:
     .. math::

        \gamma = 0.55

    see :cite:`2019:Euclid Preparation VII, eqn.32`
    """
    return Omega_m_a(cosmo, a)**cosmo.gamma


def Gf(cosmo, a):
    r"""
    FastPM growth factor function

    Parameters
    ----------
    cosmo: dict
      Cosmology dictionary.

    a : array_like
       Scale factor.

    Returns
    -------
    Scalar float Tensor : FastPM growth factor function.

    Notes
    -----

    The expression for :math:`Gf(a)` is:

    .. math::
        Gf(a)=D'_{1norm}*a**3*E(a)
    """
    f1 = growth_rate(cosmo, a)
    g1 = growth_factor(cosmo, a)
    D1f = f1 * g1 / a
    return D1f * np.power(a, 3) * np.power(Esqr(cosmo, a), 0.5)


def Gf2(cosmo, a):
    r""" FastPM second order growth factor function

    Parameters
    ----------
    cosmo: dict
      Cosmology dictionary.

    a : array_like
       Scale factor.

    Returns
    -------
    Scalar float Tensor : FastPM second order growth factor function.

    Notes
    -----

    The expression for :math:`Gf_2(a)` is:

    .. math::
        Gf_2(a)=D'_{2norm}*a**3*E(a)
    """
    f2 = growth_rate_second(cosmo, a)
    g2 = growth_factor_second(cosmo, a)
    D2f = f2 * g2 / a
    return D2f * np.power(a, 3) * np.power(Esqr(cosmo, a), 0.5)


def dGfa(cosmo, a):
    r""" Derivative of Gf against a

    Parameters
    ----------
    cosmo: dict
       Cosmology dictionary.

    a : array_like
       Scale factor.

    Returns
    -------
    Scalar float Tensor : the derivative of Gf against a.

    Notes
    -----

    The expression for :math:`gf(a)` is:

    .. math::
        gf(a)=\frac{dGF}{da}= D^{''}_1 * a ** 3 *E(a) +D'_{1norm}*a ** 3 * E'(a)
                +   3 * a ** 2 * E(a)*D'_{1norm}

    """
    f1 = growth_rate(cosmo, a)
    g1 = growth_factor(cosmo, 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)


def dGf2a(cosmo, a):
    r""" Derivative of Gf2 against a

    Parameters
    ----------
    cosmo: dict
      Cosmology dictionary.

    a : array_like
       Scale factor.

    Returns
    -------
    Scalar float Tensor : the derivative of Gf2 against a.

    Notes
    -----

    The expression for :math:`gf2(a)` is:

    .. math::
        gf_2(a)=\frac{dGF_2}{da}= D^{''}_2 * a ** 3 *E(a) +D'_{2norm}*a ** 3 * E'(a)
                +   3 * a ** 2 * E(a)*D'_{2norm}

    """
    f2 = growth_rate_second(cosmo, a)
    g2 = growth_factor_second(cosmo, 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_a = E(cosmo, a)
    return (f2p * a**3 * E_a + D2f * a**3 * dEa(cosmo, a) +
            3 * a**2 * E_a * D2f)