python 两种形式的恒星群IMF。

Posted 2021-05-09

from __future__ import division
import numpy as np

NBINS = 1000


class TappedPower:
    """
    alpha: slope for m->0
    beta: slope for m-> inf, usually beta <0. It must satisfy alpha > beta
    mstar: turning point mass
    flex: turning region size, larger flex, smaller turning radius
    mmin: min mass in consideration
    mmax: max mass in consideration
    Examples
    --------
    p = TappedPower(1, -2)
    m = np.logspace(-1, 2, 100)
    loglog(m, p.pdf(m))
    loglog(m, p.cdf(m))
    """

    def __init__(self, alpha, beta, mstar=1, mmin=0.1, mmax=100):
        self.__dict__.update(alpha=alpha, beta=beta, mstar=mstar, mmin=mmin, mmax=mmax)
        self.amp = 1
        self.amp = 1 / self.cdf(mmax)

    def pdf(self, mass):
        amp, alpha, beta, mstar = self.amp, self.alpha, self.beta, self.mstar

        return amp * mass**beta * (1 - np.exp(-(mass / mstar)**(alpha - beta)))

    def cdf(self, mass):
        mmin, mmax = self.mmin, self.mmax

        lgx, dlgx = np.linspace(np.log10(mmin), np.log10(mmax), NBINS, retstep=True)
        x = 10**lgx
        dx = x * dlgx * np.log(10)
        y = (self.pdf(x) * dx).cumsum()

        return np.interp(mass, x, y, left=np.nan, right=np.nan)

    def rvs(self, n):
        mmin, mmax = self.mmin, self.mmax

        lgx, dlgx = np.linspace(np.log10(mmin), np.log10(mmax), NBINS, retstep=True)
        x = 10**lgx
        dx = x * dlgx * np.log(10)
        y = (self.pdf(x) * dx).cumsum()

        a = np.random.rand(n)
        return np.interp(a, y, x)


class TwoPower:
    """
    alpha: slope for m->0
    beta: slope for m-> inf, usually beta <0. It must satisfy alpha > beta
    mstar: turning point mass
    flex: turning region size, larger flex, smaller turning radius
    mmin: min mass in consideration
    mmax: max mass in consideration
    Examples
    --------
    p = TwoPower(1, -2)
    m = np.logspace(-1, 2, 100)
    loglog(m, p.pdf(m))
    loglog(m, p.cdf(m))
    """

    def __init__(self, alpha, beta, mstar=1, flex=4, mmin=0.1, mmax=100):
        self.__dict__.update(alpha=alpha, beta=beta, mstar=mstar, flex=flex, mmin=mmin, mmax=mmax)
        self.amp = 1
        self.amp = 1 / self.cdf(mmax)

    def pdf(self, mass):
        amp, alpha, beta, flex, mstar = self.amp, self.alpha, self.beta, self.flex, self.mstar

        return amp * mass**alpha * (1 + (mass / mstar)**flex)**((beta - alpha) / flex)

    def cdf(self, mass):
        mmin, mmax = self.mmin, self.mmax

        lgx, dlgx = np.linspace(np.log10(mmin), np.log10(mmax), NBINS, retstep=True)
        x = 10**lgx
        dx = x * dlgx * np.log(10)
        y = (self.pdf(x) * dx).cumsum()

        return np.interp(mass, x, y, left=0)

        # from scipy.special import hyp2f1
        # x, a, b, c, d = mass, alpha, beta, flex, mstar
        # cdf_unorm = (x**(1 + a) * hyp2f1((1 + a) / c, (a - b) / c, (1 + a + c) / c, -(x / d)**c)) / (1 + a)

    def rvs(self, n):
        mmin, mmax = self.mmin, self.mmax

        lgx, dlgx = np.linspace(np.log10(mmin), np.log10(mmax), NBINS, retstep=True)
        x = 10**lgx
        dx = x * dlgx * np.log(10)
        y = (self.pdf(x) * dx).cumsum()

        a = np.random.rand(n)
        return np.interp(a, y, x)