Source code for hmf.mass_function.integrate_hmf

"""A supporting module with a routine to integrate the differential hmf in a robust manner."""

import numpy as np
import scipy.integrate as intg
from scipy.interpolate import InterpolatedUnivariateSpline as Spline


class NaNError(Exception):
    """Integrator hit a NaN."""




def hmf_integral_gtm(m, dndm, mass_density=False):
    """
    Cumulatively integrate dn/dm.

    Parameters
    ----------
    M : array_like
        Array of masses.
    dndm : array_like
        Array of dn/dm (corresponding to M)
    mass_density : bool, `False`
        Whether to calculate mass density (or number density).

    Returns
    -------
    ngtm : array_like
        Cumulative integral of dndm.

    Examples
    --------
    Using a simple power-law mass function:

    >>> import numpy as np
    >>> m = np.logspace(10,18,500)
    >>> dndm = m**-2
    >>> ngtm = hmf_integral_gtm(m,dndm)
    >>> np.allclose(ngtm,1/m) #1/m is the analytic integral to infinity.
    True

    The function always integrates to m=1e18, and extrapolates with a spline
    if data not provided:

    >>> m = np.logspace(10,12,500)
    >>> dndm = m**-2
    >>> ngtm = hmf_integral_gtm(m,dndm)
    >>> np.allclose(ngtm,1/m) #1/m is the analytic integral to infinity.
    True

    """
    n = len(m)

    # Eliminate NaN's
    mask = np.isfinite(dndm)
    m = m[mask]
    dndm = dndm[mask]
    dndlnm = m * dndm

    if len(m) < 4:
        raise NaNError(
            f"There are too few real numbers in dndm: len(dndm) = {n}, #NaN's = {n - len(m)}"
        )

    # Calculate the mass function (and its integral) from the highest M up to 10**18
    if m[-1] < m[0] * 10**18 / m[3]:
        m_upper = np.arange(np.log(m[-1]), np.log(10**18), np.log(m[1]) - np.log(m[0]))
        mf_func = Spline(np.log(m), np.log(dndlnm), k=1)
        mf = mf_func(m_upper)

        if not mass_density:
            int_upper = intg.simpson(np.exp(mf), dx=m_upper[2] - m_upper[1])
        else:
            int_upper = intg.simpson(np.exp(m_upper + mf), dx=m_upper[2] - m_upper[1])
    else:
        int_upper = 0

    # Calculate the cumulative integral (backwards) of [m*]dndlnm
    if not mass_density:
        ngtm = np.concatenate(
            (
                intg.cumulative_trapezoid(dndlnm[::-1], dx=np.log(m[1] / m[0]))[::-1],
                np.zeros(1),
            )
        )
    else:
        ngtm = np.concatenate(
            (
                intg.cumulative_trapezoid(m[::-1] * dndlnm[::-1], dx=np.log(m[1]) - np.log(m[0]))[
                    ::-1
                ],
                np.zeros(1),
            )
        )

    return ngtm + int_upper