Initial Mass Function for young star clusters and disk field stars

Question

I hope you can help with improving my code. I am defining a function which draws out elements of a certain mass, one by one, from a distribution constrained by the function imf(), until I have used up all the mass given to the function. The code takes a very long time between 1 minute to 45 minutes depending on the input mass. I am wondering if there is any way to make this code more effective? In the code there are certain parameters that give trivial answers such as log10(mnorm) this was done to ensure I could in the future alter the parameters. The focus of my problem is the while loop and how it draws from the distribution given by imf(), I have identified that this part is the root course of the long performance time for the code. Any help would be greatly appreciated.

class Mod_MyFunctions:

    def __init__(self):
        pass

    def imf(self, x, imf_type):

        # Chabrier (2003) IMF for young clusters plus disk stars: lognorm and power-law tail
        mnorm = 1.0
        A1 = 0.158
        mc = 0.079
        sigma = 0.69
        A2 = 4.43e-2
        x0 = -1.3

        if imf_type == 0:
            ml = numpy.asarray((x <= log10(mnorm)).nonzero())[0]
            mh = numpy.asarray((x > log10(mnorm)).nonzero())[0]
            y = numpy.zeros(len(x))
            for i in ml: y[i] = A1 * exp(-(x[i] - log10(mc))**2/2./sigma**2)
            for i in mh: y[i] = A2 * (10.**x[i])**(x0-1)
            return y

    def mass_dist(self,
        mmin=0.01,
        mmax=100,
        Mcm=10000,
        imf_type=0,
        SFE=0.03):

        result = []
        while sum(10**(np.array(result))) < SFE*Mcm:
            x=numpy.random.uniform(log10(mmin), log10(mmax),size=1)
            y=numpy.random.uniform(0, 1, size=1)
            result.extend(x[numpy.where(y < myf.imf(x, imf_type))])

        md=numpy.array(result)
        return 10**md, len(md)

Answer 1

Class

You use classes here without any benefit. Your methods don't use the self argument. If you use the class as a namespace, I would suggest you use a module.

np

Convention is to do import numpy as np.

You also probably did from numpy import log10, exp. I would not import those independently, just do np.exp.

vectorise

You can use np.where to select between the 2 formula's. This allows you to vectorise imf

def imf_numpy(x, imf_type):

    # Chabrier (2003) IMF for young clusters plus disk stars: lognorm and power-law tail
    mnorm = 1.0
    A1 = 0.158
    mc = 0.079
    sigma = 0.69
    A2 = 4.43e-2
    x0 = -1.3

    if imf_type == 0:
        a1 = A1 * np.exp(-((x - np.log10(mc)) ** 2) / 2.0 / sigma ** 2)
        a2 = 2 * (10.0 ** x) ** (x0 - 1)
        return np.where(x <= np.log10(mnorm), a1, a2)

I gave them a1 and a2 variable names, but I don't have any domain knowledge. If in the literature these get assigned other names, use these.

In the mass_dist, you can vectorise a lot.

By limiting your x and y to size=1, you don't get any particular benefit from using numpy. I would take larger jumps and then select how far you need them. You also use a list and extend. I would stay in numpy-land, and use hstack. You can then limit the result to SFE * Mcm:

I would keep result as the 10** already. this makes the rest easier to comprehend.

def mass_dist_numpy(mmin=0.01, mmax=100, Mcm=10000, imf_type=0, SFE=0.03):
    mmin_log = numpy.log10(mmin)
    mmax_log = numpy.log10(mmax)

    chunksize = 10
    result = np.array([], dtype=np.float64)
    while result.sum() < SFE * Mcm:
        x = np.random.uniform(mmin_log, mmax_log, size=chunksize)
        y = np.random.uniform(0, 1, size=chunksize)
        result = np.hstack((result, 10 ** x[y < imf_numpy(x, imf_type)]))

    return result[result.cumsum() < SFE * Mcm]

You can experiment with different chunk sizes, depending on the relative costs of calculating the imf, concatenating the results, the native python loop. You could pass it on as parameter and do a sensitivity analysis on it.

Answer 2

Improvise means something different then you think :)

log10(mnorm) seems to be needed to compute only once. And you can do it in your head. It is 0.

Same for log10(mc) (well not in your head this one :)).

imf_type seems useless when not zero. I am not pythonist, so, what does imf() return if imf_type is nonzero?

sigma**2 can also be computed once.

log10(mmin) and log10(mmax) can be computed once per mass_dist call.

myf.imf(x, imf_type) can be computed once per call to extend (or mass_dist, im not sure what that statement means, but I am almost sure that imf call can be moved at least one level up). This one is probably the biggest performance killer.

And there're probably more instances of this (anti)pattern.

As for the algorithm itself, I leave that to others :)