Alternative implementation of the erlang distribution

erlang.py

import numpy as np
from scipy.stats.distributions import gamma_gen


class erlang_gen(gamma_gen):

    def _argcheck(self, a):
        good = (a > 0) & (np.floor(a) == a)
        if not np.all(good):
            # This is more "aggressive" than usual, but I want, for
            # example, `erlang.mean(1.5)` to raise an error instead of
            # returning nan.
            raise ValueError('The shape parameter must be an integer.')
        return good

    def _fitstart(self, data):
        # Override _fitstart and just assign a = 1.  The value
        # doesn't matter (because the shape parameter is always fixed
        # when fitting an erlang distribution), but it must be an integer.
        a = 1
        return super(gamma_gen, self)._fitstart(data, args=(a,))

    def fit(self, data, *args, **kwds):
        if not 'f0' in kwds:
            raise ValueError('The shape parameter must be specified using f0=<integer>')
        f0 = kwds['f0']
        if f0 != np.floor(f0):
            raise ValueError('The shape parameter must be an integer')
        result = super(erlang_gen, self).fit(data, *args, **kwds)
        return result

erlang_gen.fit.__func__.__doc__ = (gamma_gen.fit.__doc__ +
"""The `fit` method of the `erlang` distribution *requires* a fixed
        shape parameter, specified with `f0=<integer>`.
""")
erlang = erlang_gen(a=0.0, name='erlang', shapes='a')