<commit_before>"""
Various codes to work with the initial mass function
"""
import numpy as np
import types # I use typechecking.  Is there a better way to do this?  (see inverse_imf below)
import scipy.integrate

class MassFunction(object):
    """
    Generic Mass Function class
    """

    def dndm(self, m, **kwargs):
        """
        The differential form of the mass function, d N(M) / dM
        """
        return self(m, integral_form=False, **kwargs)

    def n_of_m(self, m, **kwargs):
        """
        The integral form of the mass function, N(M)
        """
        return self(m, integral_form=True, **kwargs)

    def integrate(self, mlow, mhigh, **kwargs):
        """
        Integrate the mass function over some range
        """
        return scipy.integrate.quad(self, mlow, mhigh, **kwargs)

class Salpeter(MassFunction):

    def __init__(self, alpha=2.35):
        """
        Create a default Salpeter mass function, i.e. a power-law mass function 
        the Salpeter 1955 IMF: dn/dm ~ m^-2.35
        """
        self.alpha = alpha

    def __call__(self, m, integral_form=False):
        if integral_form:
            return m**(-(self.alpha - 1))
        else:
            return m**(-self.alpha)


# three codes for dn/dlog(m)
salpeter = Salpeter()

class BrokenPowerLaw(MassFunction):
    def __init__(self, breaks, mmin, mmax):
        self.breaks = breaks
        self.normalization = self.integrate(mmin, mmax)[0]

    def __call__(self, m, integral_form=False):
        zeta = 0
        b_low = 0
        alp_low = 0
        for ii,b in enumerate(self.breaks):
            if integral_form:
                alp = self.breaks[b] - 1
            else:
                alp = self.breaks[b]
            if b == 'last':
                zeta += m**(-alp) * (b_low**(-alp+alp_low)) * (m>b_low)
            else:
                mask = ((m<b)*(m>b_low))
                zeta += m**(-alp) * (b**(-alp+alp_low)) *mask
                alp_low = alp
                b_low = b

        if hasattr(self,'normalization'):
            return zeta/self.normalization
        else:
            return zeta
            
#kroupa = BrokenPowerLaw(breaks={0.08:-0.3, 0.5:1.3, 'last':2.3},mmin=0.03,mmax=120)

class Kroupa(MassFunction):
    def __init__(self, mmin=0.03):
        """
        """
        self.mmin = mmin

    def __call__(self, m, p1=0.3, p2=1.3, p3=2.3, break1=0.08, break2=0.5, integral_form=False):
        """
        Kroupa 2001 IMF (http://arxiv.org/abs/astro-ph/0009005, http://adsabs.harvard.edu/abs/2001MNRAS.322..231K)
        """

        m = np.array(m)

        binv = ((break1**(-(p1-1)) - self.mmin**(-(p1-1)))/(1-p1) +
                (break2**(-(p2-1)) - break1**(-(p2-1))) * (break1**(p2-p1))/(1-p2) +
                (- break2**(-(p3-1))) * (break1**(p2-p1)) * (break2**(p3-p2))/(1-p3))
        b = 1./binv
        c = b * break1**(p2-p1)
        d = c * break2**(p3-p2)

        zeta = (b*(m**(-(p1))) * (m<break1) +
                c*(m**(-(p2))) * (m>=break1) * (m<break2) +
                d*(m**(-(p3))) * (m>=break2))

        if integral_form:
            return zeta * m
        else:
            return zeta

kroupa = Kroupa()


def chabrier(m, integral=False):
    """
    Chabrier 2003 IMF
    http://adsabs.harvard.edu/abs/2003PASP..115..763C
    (only valid for m < 1 msun)

    not sure which of these to use...

    integral is NOT IMPLEMENTED
    """
    if integral: print "Chabrier integral NOT IMPLEMENTED"
    # This system MF can be parameterized by the same type of lognormal form as
    # the single MF (eq. [17]), with the same normalization at 1 Msun, with the
    # coefficients (Chabrier 2003)
    return 0.86 * np.exp(-1*(np.log10(m)-np.log10(0.22))**2/(2*0.57**2))
    # This analytic form for the disk MF for single objects below 1 Msun, within these uncertainties, is given by the following lognormal form (Chabrier 2003):
    return 0.158 * np.exp(-1*(np.log10(m)-np.log10(0.08))**2/(2*0.69**2))

def schechter(m,A=1,beta=2,m0=100, integral=False):
    """
    A Schechter function with arbitrary defaults
    (integral may not be correct - exponent hasn't been dealt with at all)
<commit_msg>split out lognormal from chabrier<commit_after>"""
Various codes to work with the initial mass function
"""
import numpy as np
import types # I use typechecking.  Is there a better way to do this?  (see inverse_imf below)
import scipy.integrate

class MassFunction(object):
    """
    Generic Mass Function class
    """

    def dndm(self, m, **kwargs):
        """
        The differential form of the mass function, d N(M) / dM
        """
        return self(m, integral_form=False, **kwargs)

    def n_of_m(self, m, **kwargs):
        """
        The integral form of the mass function, N(M)
        """
        return self(m, integral_form=True, **kwargs)

    def integrate(self, mlow, mhigh, **kwargs):
        """
        Integrate the mass function over some range
        """
        return scipy.integrate.quad(self, mlow, mhigh, **kwargs)

class Salpeter(MassFunction):

    def __init__(self, alpha=2.35):
        """
        Create a default Salpeter mass function, i.e. a power-law mass function 
        the Salpeter 1955 IMF: dn/dm ~ m^-2.35
        """
        self.alpha = alpha

    def __call__(self, m, integral_form=False):
        if integral_form:
            return m**(-(self.alpha - 1))
        else:
            return m**(-self.alpha)


# three codes for dn/dlog(m)
salpeter = Salpeter()

class BrokenPowerLaw(MassFunction):
    def __init__(self, breaks, mmin, mmax):
        self.breaks = breaks
        self.normalization = self.integrate(mmin, mmax)[0]

    def __call__(self, m, integral_form=False):
        zeta = 0
        b_low = 0
        alp_low = 0
        for ii,b in enumerate(self.breaks):
            if integral_form:
                alp = self.breaks[b] - 1
            else:
                alp = self.breaks[b]
            if b == 'last':
                zeta += m**(-alp) * (b_low**(-alp+alp_low)) * (m>b_low)
            else:
                mask = ((m<b)*(m>b_low))
                zeta += m**(-alp) * (b**(-alp+alp_low)) *mask
                alp_low = alp
                b_low = b

        if hasattr(self,'normalization'):
            return zeta/self.normalization
        else:
            return zeta
            
#kroupa = BrokenPowerLaw(breaks={0.08:-0.3, 0.5:1.3, 'last':2.3},mmin=0.03,mmax=120)

class Kroupa(MassFunction):
    def __init__(self, mmin=0.03):
        """
        """
        self.mmin = mmin

    def __call__(self, m, p1=0.3, p2=1.3, p3=2.3, break1=0.08, break2=0.5, integral_form=False):
        """
        Kroupa 2001 IMF (http://arxiv.org/abs/astro-ph/0009005, http://adsabs.harvard.edu/abs/2001MNRAS.322..231K)
        """

        m = np.array(m)

        binv = ((break1**(-(p1-1)) - self.mmin**(-(p1-1)))/(1-p1) +
                (break2**(-(p2-1)) - break1**(-(p2-1))) * (break1**(p2-p1))/(1-p2) +
                (- break2**(-(p3-1))) * (break1**(p2-p1)) * (break2**(p3-p2))/(1-p3))
        b = 1./binv
        c = b * break1**(p2-p1)
        d = c * break2**(p3-p2)

        zeta = (b*(m**(-(p1))) * (m<break1) +
                c*(m**(-(p2))) * (m>=break1) * (m<break2) +
                d*(m**(-(p3))) * (m>=break2))

        if integral_form:
            return zeta * m
        else:
            return zeta

kroupa = Kroupa()

def lognormal(m, offset=0.22, width=0.57, scale=0.86):
    """
    A lognormal IMF.  The default parameters correspond to the Chabrier IMF
    """
    return scale * np.exp(-1*(np.log10(m)-np.log10(offset))**2/(2*width**2))

def chabrier(m, integral=False):
    """
    Chabrier 2003 IMF
    http://adsabs.harvard.edu/abs/2003PASP..115..763C
    (only valid for m < 1 msun)

    not sure which of these to use...

    integral is NOT IMPLEMENTED
    """
    if integral:
        raise NotImplementedError("Chabrier integral NOT IMPLEMENTED")
    # This system MF can be parameterized by the same type of lognormal form as
    # the single MF (eq. [17]), with the same normalization at 1 Msun, with the
    # coefficients (Chabrier 2003)
    return lognormal(m)
    #return 0.86 * np.exp(-1*(np.log10(m)-np.log10(0.22))**2/(2*0.57**2))
    # This analytic form for the disk MF for single objects below 1 Msun, within these uncertainties, is given by the following lognormal form (Chabrier 2003):
    #return 0.158 * np.exp(-1*(np.log10(m)-np.log10(0.08))**2/(2*0.69**2))

def schechter(m,A=1,beta=2,m0=100, integral=False):
    """
    A Schechter function with arbitrary defaults
    (integral may not be correct - exponent hasn't been dealt with at all)

{'fraction_non_alphanumeric': 0.15299411561093804, 'fraction_numerical': 0.04442136840890735, 'mean_word_length': 2.1360347322720696, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2909274', 'n_tokens_mistral': 3264, 'n_tokens_neox': 2964, 'n_words': 757}