Add Gamma, Negative Binomial, and Generalized Poisson Distribution

[dachengx]

The comparisons between the implemented functions in this PR and other packages like scipy are here:

To reproduce:

from xenonnt_plot_style import XENONPlotStyle as xps
xps.use('xenonnt')


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import matplotlib.cm as cm
from scipy import stats, special
import appletree as apt
from appletree.randgen import gamma, negative_binomial, generalized_poisson


key = apt.randgen.get_key(seed=7)
batch_size = int(1e6)


alpha = 1.5
beta = 0.4


fig, ax = plt.subplots(1, 1)

x = np.linspace(0, 15, 101)

key, rvs = gamma(key, alpha, beta, shape=(batch_size,))
h = np.histogram(rvs, bins=x)
ax.hist(
    rvs,
    bins=x,
    histtype='step',
    density=True,
    label='gamma',
)

ax.hist(
    stats.gamma.rvs(alpha, size=batch_size) / beta,
    bins=x,
    histtype='step',
    density=True,
    label='scipy.stats.gamma',
)

ax.legend()
ax.set_xlabel('x')
ax.set_ylabel('PDF')

plt.show()


n = 2.5
p = 0.3


key, rvs = negative_binomial(key, p, n, shape=(batch_size,))

fig, ax = plt.subplots(1, 1)

x = np.arange(25)

b, c = np.unique(rvs, return_counts=True)
m = b <= x.max()
ax.plot(
    b[m], c[m] / c[m].sum(),
    marker='o',
    label='negative_binomial',
)

ax.plot(
    x, stats.nbinom.pmf(x, n, p),
    marker='+',
    label='scipy.stats.nbinom',
)

ax.legend()
ax.set_xlabel('x')
ax.set_ylabel('PMF')

plt.show()


lam = 1.5
eta = 0.5


def pmf(k, lam, eta):
    mu = lam * (1 - eta)
    offset = mu + eta * k
    return mu * (offset ** (k - 1)) / (special.gamma(k + 1) * np.exp(offset))

def logpmf(k, lam, eta):
    mu = lam * (1 - eta)
    offset = mu + eta * k
    return np.log(mu) + (k - 1) * np.log(offset) - special.gammaln(k + 1) - offset


fig, ax = plt.subplots(1, 1)

x = np.arange(10)

eta = 0.5
key, rvs = generalized_poisson(key, lam, eta, shape=(batch_size,))
b, c = np.unique(rvs, return_counts=True)
m = b <= x.max()
ax.plot(
    b[m], c[m] / c[m].sum(),
    marker='o',
    label=f'generalized_poisson, eta={eta:.1f}',
)
ax.plot(
    x, pmf(x, lam, eta),
    marker='+',
    label='analytical GPD',
)

eta = 0.0
key, rvs = generalized_poisson(key, lam, eta, shape=(batch_size,))
b, c = np.unique(rvs, return_counts=True)
m = b <= x.max()
ax.plot(
    b[m], c[m] / c[m].sum(),
    marker='o',
    label=f'generalized_poisson, eta={eta:.1f}',
)
ax.plot(
    x, stats.poisson.pmf(x, lam),
    marker='+',
    label='scipy.stats.poisson',
)

ax.legend(ncol=1)
ax.set_xlabel('x')
ax.set_ylabel('PMF')

plt.show()

[coveralls]

Pull Request Test Coverage Report for Build 8132432670

Details

• 9 of 31 (29.03%) changed or added relevant lines in 1 file are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage decreased (-0.7%) to 83.291%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
appletree/randgen.py	9	31	29.03%

Totals
Change from base Build 7932308533:	-0.7%
Covered Lines:	1984
Relevant Lines:	2382

---

💛 - Coveralls

[zihaoxu98]

Thanks. Looks all good to me