How to discretize a gamma or lognormal distribution, for modelling rates over sites in a phylogenetic model

discretize.py

import scipy
import scipy.stats as ss
import numpy as np

def discretize(alpha, ncat, dist=ss.gamma):
    if dist == ss.gamma:
        dist = dist(alpha, scale=1 / alpha)
    elif dist == ss.lognorm:
        dist = dist(s=alpha, scale=np.exp(0.5 * alpha**2))
    quantiles = dist.ppf(np.arange(0, ncat) / ncat)
    rates = np.zeros(ncat, dtype=np.double)
    for i in range(ncat-1):
        rates[i] = ncat * scipy.integrate.quad(lambda x: x * dist.pdf(x), 
                                               quantiles[i], quantiles[i+1])[0]
    rates[ncat-1] = ncat * scipy.integrate.quad(lambda x: x * dist.pdf(x), 
                                                quantiles[ncat-1], np.inf)[0]
    return rates