dimanche 12 février 2017

Wrapping function to use different paremetres in python

I want to generate gamma distributed random variables with scipy.stats.gamma.rvs

However, the API accepts shape and scale parameter as input and I want to use mean and standard deviation as input.

I could write a wrapper, that would get mean and std as input, computed scale and shape and than used gamma.rvs to generate the random variable. But I expect many calls of this function, so I would prefer if I didn't have to do the computation each time.

I came up with the following solution:

from scipy.stats import gamma


def gamma_rvs_generator(mu, sigma):
    var = sigma ** 2
    scale = var / mu
    a = mu / scale
    return lambda: gamma.rvs(a=a, scale=scale)


if __name__ == '__main__':
    gen = gamma_rvs_generator(10, 10)
    print(gen())

Which works exactly just as I want it to work. But I wanted to ask, if there were any better / more pythonic solutions to this.




Aucun commentaire:

Enregistrer un commentaire