File: C:/Users/fred/anaconda3/Lib/site-packages/statsmodels/discrete/tests/test_truncated_model.py
import warnings
import numpy as np
from numpy.testing import assert_allclose, assert_equal
from statsmodels import datasets
from statsmodels.tools.tools import add_constant
from statsmodels.tools.testing import Holder
from statsmodels.tools.sm_exceptions import (
ConvergenceWarning,
)
from statsmodels.distributions.discrete import (
truncatedpoisson,
truncatednegbin,
)
from statsmodels.discrete.truncated_model import (
TruncatedLFPoisson,
TruncatedLFNegativeBinomialP,
HurdleCountModel,
)
from statsmodels.sandbox.regression.tests.test_gmm_poisson import DATA
from .results.results_discrete import RandHIE
from .results import results_truncated as results_t
from .results import results_truncated_st as results_ts
class CheckResults:
def test_params(self):
assert_allclose(self.res1.params, self.res2.params,
atol=1e-5, rtol=1e-5)
def test_llf(self):
assert_allclose(self.res1.llf, self.res2.llf, atol=1e-5, rtol=1e-7)
def test_conf_int(self):
assert_allclose(self.res1.conf_int(), self.res2.conf_int,
atol=1e-3, rtol=1e-5)
def test_bse(self):
assert_allclose(self.res1.bse, self.res2.bse, atol=1e-3)
def test_aic(self):
assert_allclose(self.res1.aic, self.res2.aic, atol=1e-2, rtol=1e-12)
def test_bic(self):
assert_allclose(self.res1.bic, self.res2.bic, atol=1e-2, rtol=1e-12)
def test_fit_regularized(self):
model = self.res1.model
alpha = np.ones(len(self.res1.params))
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=ConvergenceWarning)
# This does not catch all Convergence warnings, why?
res_reg = model.fit_regularized(alpha=alpha*0.01, disp=0)
assert_allclose(res_reg.params, self.res1.params,
rtol=1e-3, atol=5e-3)
assert_allclose(res_reg.bse, self.res1.bse,
rtol=1e-3, atol=5e-3)
class TestTruncatedLFPoissonModel(CheckResults):
@classmethod
def setup_class(cls):
data = datasets.randhie.load()
exog = add_constant(np.asarray(data.exog)[:, :4], prepend=False)
mod = TruncatedLFPoisson(data.endog, exog, truncation=5)
cls.res1 = mod.fit(method="newton", maxiter=500)
res2 = RandHIE()
res2.truncated_poisson()
cls.res2 = res2
class TestZeroTruncatedLFPoissonModel(CheckResults):
@classmethod
def setup_class(cls):
data = datasets.randhie.load()
exog = add_constant(np.asarray(data.exog)[:, :4], prepend=False)
mod = TruncatedLFPoisson(data.endog, exog, truncation=0)
cls.res1 = mod.fit(maxiter=500)
res2 = RandHIE()
res2.zero_truncated_poisson()
cls.res2 = res2
class TestZeroTruncatedNBPModel(CheckResults):
@classmethod
def setup_class(cls):
data = datasets.randhie.load()
exog = add_constant(np.asarray(data.exog)[:, :3], prepend=False)
mod = TruncatedLFNegativeBinomialP(data.endog, exog, truncation=0)
cls.res1 = mod.fit(maxiter=500)
res2 = RandHIE()
res2.zero_truncted_nbp()
cls.res2 = res2
def test_conf_int(self):
pass
class TestTruncatedLFPoisson_predict:
@classmethod
def setup_class(cls):
cls.expected_params = [1, 0.5]
np.random.seed(123)
nobs = 200
exog = np.ones((nobs, 2))
exog[:nobs//2, 1] = 2
mu_true = exog.dot(cls.expected_params)
cls.endog = truncatedpoisson.rvs(mu_true, 0, size=mu_true.shape)
model = TruncatedLFPoisson(cls.endog, exog, truncation=0)
cls.res = model.fit(method='bfgs', maxiter=5000)
def test_mean(self):
assert_allclose(self.res.predict().mean(), self.endog.mean(),
atol=2e-1, rtol=2e-1)
def test_var(self):
v = self.res.predict(which="var").mean()
assert_allclose(v, self.endog.var(), atol=2e-1, rtol=2e-1)
return
assert_allclose((self.res.predict().mean() *
self.res._dispersion_factor.mean()),
self.endog.var(), atol=5e-2, rtol=5e-2)
def test_predict_prob(self):
res = self.res
pr = res.predict(which='prob')
pr2 = truncatedpoisson.pmf(
np.arange(8), res.predict(which="mean-main")[:, None], 0)
assert_allclose(pr, pr2, rtol=1e-10, atol=1e-10)
class TestTruncatedNBP_predict:
@classmethod
def setup_class(cls):
cls.expected_params = [1, 0.5, 0.5]
np.random.seed(1234)
nobs = 200
exog = np.ones((nobs, 2))
exog[:nobs//2, 1] = 2
mu_true = np.exp(exog.dot(cls.expected_params[:-1]))
cls.endog = truncatednegbin.rvs(
mu_true, cls.expected_params[-1], 2, 0, size=mu_true.shape)
model = TruncatedLFNegativeBinomialP(cls.endog, exog,
truncation=0, p=2)
cls.res = model.fit(method='nm', maxiter=5000, maxfun=5000)
def test_mean(self):
assert_allclose(self.res.predict().mean(), self.endog.mean(),
atol=2e-1, rtol=2e-1)
def test_var(self):
v = self.res.predict(which="var").mean()
assert_allclose(v, self.endog.var(), atol=1e-1, rtol=1e-2)
return
assert_allclose((self.res.predict().mean() *
self.res._dispersion_factor.mean()),
self.endog.var(), atol=5e-2, rtol=5e-2)
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)
class CheckTruncatedST():
def test_basic(self):
res1 = self.res1
res2 = self.res2
assert_allclose(res1.llf, res2.ll, rtol=1e-8)
assert_allclose(res1.llnull, res2.ll_0, rtol=5e-6)
pt2 = res2.params_table
# Stata has different parameterization of alpha for negbin
k = res1.model.exog.shape[1]
assert_allclose(res1.params[:k], res2.params[:k], atol=1e-5)
assert_allclose(res1.bse[:k], pt2[:k, 1], atol=1e-5)
assert_allclose(res1.tvalues[:k], pt2[:k, 2], rtol=5e-4, atol=5e-4)
assert_allclose(res1.pvalues[:k], pt2[:k, 3], rtol=5e-4, atol=1e-7)
assert_equal(res1.df_model, res2.df_m)
assert_allclose(res1.aic, res2.icr[-2], rtol=1e-8)
assert_allclose(res1.bic, res2.icr[-1], rtol=1e-8)
nobs = res1.model.endog.shape[0]
assert_equal((res1.model.endog < 1).sum(), 0)
# df_resid not available in Stata
assert_equal(res1.df_resid, nobs - len(res1.params))
def test_predict(self):
res1 = self.res1
res2 = self.res2
# mean of untruncated distribution
rdf = res2.margins_means.table
pred = res1.get_prediction(which="mean-main", average=True)
assert_allclose(pred.predicted, rdf[0], rtol=5e-5)
assert_allclose(pred.se, rdf[1], rtol=5e-4, atol=1e-10)
ci = pred.conf_int()[0]
assert_allclose(ci[0], rdf[4], rtol=1e-5, atol=1e-10)
assert_allclose(ci[1], rdf[5], rtol=1e-5, atol=1e-10)
# mean of untruncated distribution, evaluated and exog.mean()
ex = res1.model.exog.mean(0)
rdf = res2.margins_atmeans.table
pred = res1.get_prediction(ex, which="mean-main")
assert_allclose(pred.predicted, rdf[0], rtol=5e-5)
assert_allclose(pred.se, rdf[1], rtol=5e-4, atol=1e-10)
ci = pred.conf_int()[0]
assert_allclose(ci[0], rdf[4], rtol=5e-5, atol=1e-10)
assert_allclose(ci[1], rdf[5], rtol=5e-5, atol=1e-10)
# mean of truncated distribution, E(y | y > trunc)
rdf = res2.margins_cm.table
try:
pred = res1.get_prediction(average=True)
except NotImplementedError:
# not yet implemented for truncation > 0
pred = None
if pred is not None:
assert_allclose(pred.predicted, rdf[0], rtol=5e-5)
assert_allclose(pred.se, rdf[1], rtol=1e-5, atol=1e-10)
ci = pred.conf_int()[0]
assert_allclose(ci[0], rdf[4], rtol=1e-5, atol=1e-10)
assert_allclose(ci[1], rdf[5], rtol=1e-5, atol=1e-10)
# predicted probabilites, only subset is common to reference
ex = res1.model.exog.mean(0)
rdf = res2.margins_cpr.table
start_idx = res1.model.truncation + 1
k = rdf.shape[0] + res1.model.truncation
pred = res1.get_prediction(which="prob", average=True)
assert_allclose(pred.predicted[start_idx:k], rdf[:-1, 0], rtol=5e-5)
assert_allclose(pred.se[start_idx:k], rdf[:-1, 1],
rtol=5e-4, atol=1e-10)
ci = pred.conf_int()[start_idx:k]
assert_allclose(ci[:, 0], rdf[:-1, 4], rtol=5e-5, atol=1e-10)
assert_allclose(ci[:, 1], rdf[:-1, 5], rtol=5e-5, atol=1e-10)
# untruncated predicted probabilites, subset is common to reference
ex = res1.model.exog.mean(0)
rdf = res2.margins_pr.table
k = rdf.shape[0] - 1
pred = res1.get_prediction(which="prob-base", average=True)
assert_allclose(pred.predicted[:k], rdf[:-1, 0], rtol=5e-5)
assert_allclose(pred.se[:k], rdf[:-1, 1],
rtol=8e-4, atol=1e-10)
ci = pred.conf_int()[:k]
assert_allclose(ci[:, 0], rdf[:-1, 4], rtol=5e-4, atol=1e-10)
assert_allclose(ci[:, 1], rdf[:-1, 5], rtol=5e-4, atol=1e-10)
class TestTruncatedLFPoissonSt(CheckTruncatedST):
# test against R pscl
@classmethod
def setup_class(cls):
endog = DATA["docvis"]
exog_names = ['aget', 'totchr', 'const']
exog = DATA[exog_names]
cls.res1 = TruncatedLFPoisson(endog, exog).fit(method="bfgs",
maxiter=300)
cls.res2 = results_ts.results_trunc_poisson
mod_offset = TruncatedLFPoisson(endog, exog, offset=DATA["aget"])
cls.res_offset = mod_offset.fit(method="bfgs", maxiter=300)
def test_offset(self):
res1 = self.res1
reso = self.res_offset
paramso = np.asarray(reso.params)
params1 = np.asarray(res1.params)
assert_allclose(paramso[1:], params1[1:], rtol=1e-8)
assert_allclose(paramso[0], params1[0] - 1, rtol=1e-8)
pred1 = res1.predict()
predo = reso.predict()
assert_allclose(predo, pred1, rtol=1e-8)
ex = res1.model.exog[:5]
offs = reso.model.offset[:5]
pred1 = res1.predict(ex, transform=False)
predo = reso.predict(ex, offset=offs, transform=False)
assert_allclose(predo, pred1, rtol=1e-8)
class TestTruncatedNegBinSt(CheckTruncatedST):
# test against R pscl
@classmethod
def setup_class(cls):
endog = DATA["docvis"]
exog_names = ['aget', 'totchr', 'const']
exog = DATA[exog_names]
cls.res1 = TruncatedLFNegativeBinomialP(endog, exog).fit(method="bfgs",
maxiter=300)
cls.res2 = results_ts.results_trunc_negbin
mod_offset = TruncatedLFNegativeBinomialP(endog, exog,
offset=DATA["aget"])
cls.res_offset = mod_offset.fit(method="bfgs", maxiter=300)
def test_offset(self):
# identical, copy of method in TestTruncatedLFPoissonSt
res1 = self.res1
reso = self.res_offset
paramso = np.asarray(reso.params)
params1 = np.asarray(res1.params)
assert_allclose(paramso[1:], params1[1:], rtol=1e-8)
assert_allclose(paramso[0], params1[0] - 1, rtol=1e-8)
pred1 = res1.predict()
predo = reso.predict()
assert_allclose(predo, pred1, rtol=1e-8)
ex = res1.model.exog[:5]
offs = reso.model.offset[:5]
pred1 = res1.predict(ex, transform=False)
predo = reso.predict(ex, offset=offs, transform=False)
assert_allclose(predo, pred1, rtol=1e-8)
class TestTruncatedLFPoisson1St(CheckTruncatedST):
# test against R pscl
@classmethod
def setup_class(cls):
endog = DATA["docvis"]
exog_names = ['aget', 'totchr', 'const']
exog = DATA[exog_names]
cls.res1 = TruncatedLFPoisson(
endog, exog, truncation=1
).fit(method="bfgs", maxiter=300)
cls.res2 = results_ts.results_trunc_poisson1
class TestTruncatedNegBin1St(CheckTruncatedST):
# test against R pscl
@classmethod
def setup_class(cls):
endog = DATA["docvis"]
exog_names = ['aget', 'totchr', 'const']
exog = DATA[exog_names]
cls.res1 = TruncatedLFNegativeBinomialP(
endog, exog, truncation=1
).fit(method="newton", maxiter=300) # "bfgs" is not close enough
cls.res2 = results_ts.results_trunc_negbin1
class TestHurdlePoissonR():
# test against R pscl
@classmethod
def setup_class(cls):
endog = DATA["docvis"]
exog_names = ['const', 'aget', 'totchr']
exog = DATA[exog_names]
cls.res1 = HurdleCountModel(endog, exog).fit(method="newton",
maxiter=300)
cls.res2 = results_t.hurdle_poisson
def test_basic(self):
res1 = self.res1
res2 = self.res2
assert_allclose(res1.llf, res2.loglik, rtol=1e-8)
pt2 = res2.params_table
assert_allclose(res1.params, pt2[:, 0], atol=1e-5)
assert_allclose(res1.bse, pt2[:, 1], atol=1e-5)
assert_allclose(res1.tvalues, pt2[:, 2], rtol=5e-4, atol=5e-4)
assert_allclose(res1.pvalues, pt2[:, 3], rtol=5e-4, atol=1e-7)
assert_equal(res1.df_resid, res2.df_residual)
assert_equal(res1.df_model, res2.df_null - res2.df_residual)
assert_allclose(res1.aic, res2.aic, rtol=1e-8)
# we have zero model first
idx = np.concatenate((np.arange(3, 6), np.arange(3)))
vcov = res2.vcov[idx[:, None], idx]
assert_allclose(np.asarray(res1.cov_params()), vcov,
rtol=1e-4, atol=1e-8)
def test_predict(self):
res1 = self.res1
res2 = self.res2
ex = res1.model.exog.mean(0, keepdims=True)
mu1 = res1.results_zero.predict(ex)
prob_zero = np.exp(-mu1)
prob_nz = 1 - prob_zero
assert_allclose(prob_nz, res2.predict_zero, rtol=5e-4, atol=5e-4)
prob_nz_ = res1.results_zero.model._prob_nonzero(mu1, res1.params[:4])
assert_allclose(prob_nz_, res2.predict_zero, rtol=5e-4, atol=5e-4)
mean_main = res1.results_count.predict(ex, which="mean-main")
assert_allclose(mean_main, res2.predict_mean_main,
rtol=5e-4, atol=5e-4)
prob_main = res1.results_count.predict(ex, which="prob")[0] * prob_nz
prob_main[0] = np.squeeze(prob_zero)
assert_allclose(prob_main[:4], res2.predict_prob, rtol=5e-4, atol=5e-4)
assert_allclose(mean_main * prob_nz, res2.predict_mean,
rtol=1e-3, atol=5e-4)
# with corresponding predict `which`
m = res1.predict(ex)
assert_allclose(m, res2.predict_mean, rtol=1e-6, atol=5e-7)
mm = res1.predict(ex, which="mean-main")
assert_allclose(mm, res2.predict_mean_main, rtol=1e-7, atol=1e-7)
mnz = res1.predict(ex, which="mean-nonzero")
assert_allclose(mnz, res2.predict_mean / (1 - res2.predict_prob[0]),
rtol=5e-7, atol=5e-7)
prob_main = res1.predict(ex, which="prob-main")
pt = res1.predict(ex, which="prob-trunc")
assert_allclose(prob_main / (1 - pt), res2.predict_zero,
rtol=5e-4, atol=5e-4)
probs = res1.predict(ex, which="prob")[0] # return is 2-dim
assert_allclose(probs[:4], res2.predict_prob, rtol=1e-5, atol=1e-6)
# check vectorized and options, consistencey and smoke
k_ex = 5
ex5 = res1.model.exog[:k_ex]
p1a = res1.predict(ex5, which="prob", y_values=np.arange(3))
p1b = res1.get_prediction(ex5, which="prob", y_values=np.arange(3))
assert_allclose(p1a, p1b.predicted, rtol=1e-10, atol=1e-10)
# TODO: two dim prediction not yet supported in frame
# assert p1b.summary_frame().shape == (4, 4)
p2a = res1.predict(which="prob", y_values=np.arange(3))
p2b = res1.get_prediction(which="prob", y_values=np.arange(3),
average=True)
assert_allclose(p2a.mean(0), p2b.predicted, rtol=1e-10, atol=1e-10)
# TODO: which="var" raises AttributeError
for which in ["mean", "mean-main", "prob-main", "prob-zero", "linear"]:
p3a = res1.predict(ex5, which=which)
p3b = res1.get_prediction(ex5, which=which)
assert_allclose(p3a, p3b.predicted, rtol=1e-10, atol=1e-10)
assert p3b.summary_frame().shape == (k_ex, 4)
# var1 = res1.predict(which="var")
resid_p1 = res1.resid_pearson[:5]
resid_p2 = np.asarray([
-1.5892397298897, -0.3239276467705, -1.5878941800178,
0.6613236544236, -0.6690997162962,
])
assert_allclose(resid_p1, resid_p2, rtol=1e-5, atol=1e-5)
class CheckHurdlePredict():
def test_basic(self):
res1 = self.res1
res2 = self.res2
assert res1.df_model == res2.df_model
# assert res1.df_null == res2.df_null # not in res1
assert res1.df_resid == res2.df_resid
assert res1.model.k_extra == res2.k_extra
assert len(res1.model.exog_names) == res2.k_params
assert res1.model.exog_names == res2.exog_names
# smoke test
res1.summary()
def test_predict(self):
res1 = self.res1
endog = res1.model.endog
exog = res1.model.exog
pred_mean = res1.predict(which="mean").mean()
assert_allclose(pred_mean, endog.mean(), rtol=1e-2)
mask_nz = endog > 0
mean_nz = endog[mask_nz].mean()
pred_mean_nz = res1.predict(which="mean-nonzero").mean()
assert_allclose(pred_mean_nz, mean_nz, rtol=0.05)
# Note: the Truncated model is based on different exog
# prediction for nonzero part is better in nonzero sample than full
pred_mean_nnz = res1.predict(exog=exog[mask_nz],
which="mean-nonzero").mean()
assert_allclose(pred_mean_nnz, mean_nz, rtol=5e-4)
pred_mean_nzm = res1.results_count.predict(which="mean").mean()
assert_allclose(pred_mean_nzm, mean_nz, rtol=5e-4)
assert_allclose(pred_mean_nzm, pred_mean_nnz, rtol=1e-4)
# check variance
pred_var = res1.predict(which="var").mean()
assert_allclose(pred_var, res1.resid.var(), rtol=0.05)
pred_var = res1.results_count.predict(which="var").mean()
assert_allclose(pred_var, res1.resid[endog > 0].var(), rtol=0.05)
# check probabilities
freq = np.bincount(endog.astype(int)) / len(endog)
pred_prob = res1.predict(which="prob").mean(0)
assert_allclose(pred_prob, freq, rtol=0.005, atol=0.01)
dia_hnb = res1.get_diagnostic()
assert_allclose(dia_hnb.probs_predicted.mean(0), pred_prob, rtol=1e-10)
try:
dia_hnb.plot_probs()
except ImportError:
pass
pred_prob0 = res1.predict(which="prob-zero").mean(0)
assert_allclose(pred_prob0, freq[0], rtol=1e-4)
assert_allclose(pred_prob0, pred_prob[0], rtol=1e-10)
class TestHurdleNegbinSimulated(CheckHurdlePredict):
@classmethod
def setup_class(cls):
nobs = 2000
exog = np.column_stack((np.ones(nobs), np.linspace(0, 3, nobs)))
y_fake = np.arange(nobs) // (nobs / 3) # need some zeros and non-zeros
# get predicted probabilities for model
mod = HurdleCountModel(y_fake, exog, dist="negbin", zerodist="negbin")
p_dgp = np.array([-0.4, 2, 0.5, 0.2, 0.5, 0.5])
probs = mod.predict(p_dgp, which="prob", y_values=np.arange(50))
cdf = probs.cumsum(1)
n = cdf.shape[0]
cdf = np.column_stack((cdf, np.ones(n)))
# simulate data,
# cooked example that doesn't have identification problems
rng = np.random.default_rng(987456348)
u = rng.random((n, 1))
endog = np.argmin(cdf < u, axis=1)
mod_hnb = HurdleCountModel(endog, exog,
dist="negbin", zerodist="negbin")
cls.res1 = mod_hnb.fit(maxiter=300)
df_null = 4
cls.res2 = Holder(
nobs=nobs,
k_params=6,
df_model=2,
df_null=df_null,
df_resid=nobs-6,
k_extra=df_null - 1,
exog_names=['zm_const', 'zm_x1', 'zm_alpha', 'const', 'x1',
'alpha'],
)