def test_pmf(self): n, p = truncatednegbin.convert_params(2, 0.5, 2) nb_logpmf = nbinom.pmf(6, n, p) / nbinom.sf(5, n, p) tnb_pmf = truncatednegbin.pmf(6, 2, 0.5, 2, 5) assert_allclose(nb_logpmf, tnb_pmf, rtol=1e-7) tnb_pmf = truncatednegbin.pmf(5, 2, 0.5, 2, 5) assert_equal(tnb_pmf, 0)
def test_predict_prob(self): res = self.res pr = res.predict(which='prob') pr2 = truncatednegbin.pmf( np.arange(29), res.predict(which="mean-main")[:, None], res.params[-1], 2, 0) assert_allclose(pr, pr2, rtol=1e-10, atol=1e-10)
def test_pmf_zero(self): n, p = truncatednegbin.convert_params(5, 0.1, 2) nb_pmf = nbinom.pmf(1, n, p) / nbinom.sf(0, n, p) tnb_pmf = truncatednegbin.pmf(1, 5, 0.1, 2, 0) assert_allclose(nb_pmf, tnb_pmf, rtol=1e-5)