Alternative scipy erlang class, take 2

erlang2.py

import warnings
import numpy as np
from scipy.stats.distributions import gamma_gen, rv_continuous


class erlang_gen(gamma_gen):
    """An Erlang continuous random variable.

    %(before_notes)s

    See Also
    --------
    gamma

    Notes
    -----
    The Erlang distribution is a special case of the Gamma distribution,
    with the shape parameter `a` an integer.  This class is simply a
    subclass of the `gamma` distribution.  It behaves the same as the
    `gamma` distribution, but it will generate a warning the first time
    a non-integer value is used for the shape parameter.

    Refer to the `gamma` distribution for further examples.

    """

    def _argcheck(self, a):
        allint = np.all(np.floor(a) == a)
        allpos = np.all(a > 0) 
        if not allint:
            # An Erlang distribution shouldn't really have a non-integer
            # shape parameter, so warn the user.
            warnings.warn('The shape parameter of the erlang distribution '
                'has been given a non-integer value.', RuntimeWarning)
        return allpos

    # Trivial override of the fit method, so we can monkey-patch its
    # docstring.
    def fit(self, data, *args, **kwds):
        return super(erlang_gen, self).fit(data, *args, **kwds)

    fit.__doc__ = (rv_continuous.fit.__doc__ +
        """
        Notes
        -----
        The Erlang distribution is generally defined to have integer values
        for the shape parameter.  This is not enforced by the `erlang` class.
        When fitting the distribution, it will generally return a non-integer
        value for the shape parameter.  By using the keyword argument
        `f0=<integer>`, the fit method can be constrained to fit the data to
        a specifc integer shape parameter. 
        """)


erlang = erlang_gen(a=0.0, name='erlang', shapes='a')