a
    \:bQ                     @   s  d Z ddlmZ ddlZddlmZmZmZ ddl	Z
ddlZddlZddlmZ ddlmZ ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZmZ ej Ze ej!e_!e ej"e_"eej"dde_"dZ#dZ$e
j%ee$ddZ&e&'e(Z&e)e&d e&d< G dd de*Z+G dd de+Z,G dd de+Z-G dd de+Z.G dd de+Z/G dd de+Z0G d d! d!e+Z1G d"d# d#e+Z2G d$d% d%e+Z3G d&d' d'e+Z4G d(d) d)e+Z5G d*d+ d+e5Z6G d,d- d-e5Z7G d.d/ d/e5Z8d0d1 Z9dS )2zM
Created on Fri May 30 16:22:29 2014

Author: Josef Perktold
License: BSD-3

    )StringION)assert_assert_allcloseassert_equal)datasetsfit_constrained)PoissonfamiliesGLM)add_constant   )results_glm_logit_constrainedresults_poisson_constrainedF)Zprependzagecat	smokes	deaths	pyears
1	1	32	52407
2	1	104	43248
3	1	206	28612
4	1	186	12663
5	1	102	5317
1	0	2	18790
2	0	12	10673
3	0	28	5710
4	0	28	2585
5	0	31	1462	)Z	delimiterpyearsZ	logpyearsc                   @   s$   e Zd Zdd Zdd Zdd ZdS )CheckPoissonConstrainedMixinc                 C   sz   | j }| j}t|d |j| j dd tt|d }|dkt|j	| j @ }t||  |j	| j |  dd d S )Nr   ư>Zrtolr   )
res1res2r   paramsidxnpZsqrtZdiagisnanbse)selfr   r   Zbse1mask r    Jlib/python3.9/site-packages/statsmodels/discrete/tests/test_constrained.py
test_basic9   s    z'CheckPoissonConstrainedMixin.test_basicc           
      C   s  t | drt | jds| jn| jj}| j}t|j|j| j dd |jdkt	|j| j @ }t|j|  |j| j |  dd |j
| jdf }t|jt	|@ }t|j|  ||  dd |j
| jdf }|jdkt	|@ }t|j|  ||  dd |j
| jd	f }|j
| jd
f }t||f}t| t	|  |t	|  dd t|j|jdd t|j|j |j|j d }	t|j|	 n
td d S )Nres1m_resultsr   r   r         g-C6
?      r   not available yet)hasattrr#   r$   r   r   r   r   r   r   r   Zparams_tableZisinftvaluespvaluesZcolumn_stackZconf_intZllfZllr   Zdf_modelZdf_mNZdf_residpytestskip)
r   r   r   r   r+   r,   Zci_lowZci_uppZciZdf_rr    r    r!   test_basic_methodB   s6    "z.CheckPoissonConstrainedMixin.test_basic_methodc                 C   sd   t | drV| j}| j}t |dr4t|j|jdd q`tr`dd l}d|j }|| n
t	
d d S )Nr#   ll_0r   r   r   z&test: ll_0 not available, llnull=%6.4Fr)   )r*   r#   r   r   Zllnullr1   DEBUGwarningswarnr.   r/   )r   r   r   r3   messager    r    r!   
test_otherl   s    

z'CheckPoissonConstrainedMixin.test_otherN)__name__
__module____qualname__r"   r0   r6   r    r    r    r!   r   7   s   	*r   c                   @   s8   e Zd Zedd Zejjdd Zejjdd Z	dS )TestPoissonConstrained1ac                 C   s   t j| _g d| _d}tj|td}t|j}t	
|}t	|j |d< d}t|j|}t||j|j|dddd| _|j||ddd	| _d S )
N   r&   r'   r(      r   r   'deaths ~ logpyears + smokes + C(agecat)datar   C(agecat)[T.4] = C(agecat)[T.5]bfgsmethoddispstart_paramsfit_kwdsrG   rD   rE   )resultsresults_noexposure_constraintr   r   r	   from_formular@   len
exog_namesr   zeroslogendogmeanpatsy
DesignInfolinear_constraintr   coefs	constantsr   r#   clsformulamodZk_varsrG   constrlcr    r    r!   setup_class   s     


z$TestPoissonConstrained1a.setup_classc                 C   s   | j  }td|jv  d S )Nlinear equality constraints)r#   summaryr   	extra_txtr   summr    r    r!   test_summary   s    
z%TestPoissonConstrained1a.test_summaryc                 C   s    | j  }td|jd v  d S )Nr_   r   )r#   summary2r   ra   rb   r    r    r!   test_summary2   s    
z&TestPoissonConstrained1a.test_summary2N)
r7   r8   r9   classmethodr^   r.   marksmokerd   rf   r    r    r    r!   r:      s   

r:   c                   @   s   e Zd Zedd ZdS )TestPoissonConstrained1bc                 C   sz   t j| _g d| _d}tj|ttd jd}d}t	|j
|}t||j|jdddd	| _|| _|j|ddd| _d S )
Nr=   r%   r&   r'   r(   r   deaths ~ smokes + C(agecat)r   r@   ZexposurerA   newtonr   rC   rH   )rJ   results_exposure_constraintr   r   r	   rL   r@   valuesrS   rT   rN   rU   r   rV   rW   r   constraintsr#   rY   rZ   r[   r\   r]   r    r    r!   r^      s"    
z$TestPoissonConstrained1b.setup_classNr7   r8   r9   rg   r^   r    r    r    r!   rj      s   rj   c                   @   s   e Zd Zedd ZdS )TestPoissonConstrained1cc                 C   s   t j| _g d| _d}tj|tttd j	d}d}t
|j|}t||j|jdddd	| _|| _|j|ddd| _d S )
Nrk   rl   r   r@   offsetrA   rn   r   rC   ro   )rJ   rp   r   r   r	   rL   r@   r   rP   rq   rS   rT   rN   rU   r   rV   rW   r   rr   r#   rs   r    r    r!   r^      s    
z$TestPoissonConstrained1c.setup_classNrt   r    r    r    r!   ru      s   ru   c                   @   s   e Zd Zedd ZdS )TestPoissonNoConstrainedc                 C   sZ   t j| _g d| _d}tj|tttd j	d}|j
ddj}|j| f| _|| _d S )Nrk   rl   r   rv   r   )rE   )rJ   Zresults_exposure_noconstraintr   r   r	   rL   r@   r   rP   rq   fitr$   r   Z
cov_paramsr   r#   )rY   rZ   r[   r   r    r    r!   r^      s    
z$TestPoissonNoConstrained.setup_classNrt   r    r    r    r!   rx      s   rx   c                   @   s   e Zd Zedd ZdS )TestPoissonConstrained2ac                 C   s   t j| _g d| _d}tj|td}t|j}t	
|}t	|j |d< d}t|j|}t||j|j|dddd| _|j||ddd	| _d S )
Nr;   r>   r?   r   %C(agecat)[T.5] - C(agecat)[T.4] = 0.5rB   rC   rF   rI   )rJ   Zresults_noexposure_constraint2r   r   r	   rL   r@   rM   rN   r   rO   rP   rQ   rR   rS   rT   rU   r   rV   rW   r   r#   rX   r    r    r!   r^      s     


z$TestPoissonConstrained2a.setup_classNrt   r    r    r    r!   rz      s   rz   c                   @   s   e Zd Zedd ZdS )TestPoissonConstrained2bc                 C   s   t j| _g d| _d}tj|ttd jd}d}t	|j
|}t||j|jdddd	| _|| _|j|d
d| jd d| _d S )Nrk   rl   r   rm   r{   rn   r   rC   ro   rB   rD   rE   rG   )rJ   results_exposure_constraint2r   r   r	   rL   r@   rq   rS   rT   rN   rU   r   rV   rW   r   rr   r#   rs   r    r    r!   r^     s"    

z$TestPoissonConstrained2b.setup_classNrt   r    r    r    r!   r|     s   r|   c                   @   s   e Zd Zedd ZdS )TestPoissonConstrained2cc                 C   s   t j| _g d| _d}tj|tttd j	d}d}t
|j|}t||j|jdddd	| _|| _|j|d
d| jd d| _d S )Nrk   rl   r   rv   r{   rn   r   rC   ro   rB   r}   )rJ   r~   r   r   r	   rL   r@   r   rP   rq   rS   rT   rN   rU   r   rV   rW   r   rr   r#   rs   r    r    r!   r^   +  s     
z$TestPoissonConstrained2c.setup_classNrt   r    r    r    r!   r   )  s   r   c                   @   s   e Zd Zedd ZdS )TestGLMPoissonConstrained1ac                 C   s   ddl m} tj| _g d| _d}tj|tt	
 d}d}t|j|}|||j|jddid	| _|| _|j|dd
| _d S )Nr   r   r;   r>   )r@   familyrA   atol绽|=ro   r   )statsmodels.base._constraintsr   rJ   rK   r   r   r   rL   r@   r   r	   rS   rT   rN   rU   rV   rW   r   rr   r#   )rY   r   rZ   r[   r\   r]   r    r    r!   r^   E  s    
z'TestGLMPoissonConstrained1a.setup_classNrt   r    r    r    r!   r   C  s   r   c                   @   s    e Zd Zedd Zdd ZdS )TestGLMPoissonConstrained1bc                 C   s   ddl m} ddlm} ddlm} tj| _g d| _	d}|j
|t| ttd jd}d	}t|j|}|||j|jd
did| _|| _|j|ddj| _d S )Nr   r   r
   r   rk   rl   r   )r@   r   rw   rA   r   r   ro   r   )r   r   statsmodels.genmodr   +statsmodels.genmod.generalized_linear_modelr   rJ   rp   r   r   rL   r@   r	   r   rP   rq   rS   rT   rN   rU   rV   rW   r   rr   r$   r#   )rY   r   r   r   rZ   r[   r\   r]   r    r    r!   r^   \  s"    
z'TestGLMPoissonConstrained1b.setup_classc                 C   s   | j }| j}d}tj|ttd jd}d}|j|| j jdddd}t|j|jd	d
 t|j	|j	dd
 |
 }t||
 dd
 t|j|dd
 t|j|dd
 t|j
dd|j
dddd
 d S )Nrl   r   rm   rA   rn   Fr   )rG   rD   Zwarn_convergencerE   g-q=r   gdy=r   T)Zlinear)r#   r   r	   rL   r@   rq   r   r   r   r   predictmufittedvalues)r   r   r   rZ   r[   r\   	predictedr    r    r!   test_compare_glm_poissons  s(    z4TestGLMPoissonConstrained1b.test_compare_glm_poissonN)r7   r8   r9   rg   r^   r   r    r    r    r!   r   Z  s   
r   c                   @   s   e Zd Zdd ZdS )CheckGLMConstrainedMixinc                 C   s   | j }| j}t|j|jd dd dd l}| . |dt t|j	|j	dd W d    n1 sf0    Y  t|j
|j
dd d S )Nr'   r   r   r   ignore)r   r#   r   ZaicZinfocritr3   catch_warningssimplefilterFutureWarningZbicZdeviance)r   r   r   r3   r    r    r!   test_glm  s    
0z!CheckGLMConstrainedMixin.test_glmN)r7   r8   r9   r   r    r    r    r!   r     s   r   c                   @   s   e Zd Zedd ZdS )TestGLMLogitConstrained1c                 C   sV   t d | _tj| _ttjtjt	
 d}d}||| _| jj\}}t|||| _d S )Nr   zx1 = 2.8)slicer   reslogitZresults_constraint1r   r   spector_datarQ   exogr   Binomialr   r#   rr   r   rY   mod1r\   Rqr    r    r!   r^     s    

z$TestGLMLogitConstrained1.setup_classNrt   r    r    r    r!   r     s   r   c                   @   sH   e Zd Zedd Zdd Zejjdd Z	ejjdd Z
d	d
 ZdS )TestGLMLogitConstrained2c                 C   sv   t d | _tj| _ttjtjt	
 d}d}|j|dd| _| jjj| jjj }}t|||ddid| _||f| _d S )Nr   x1 - x3 = 0r   r   r   ro   )r   r   r   Zresults_constraint2r   r   r   rQ   r   r   r   r   r#   rr   rV   rW   r   constraints_rqr   r    r    r!   r^     s    

z$TestGLMLogitConstrained2.setup_classc                 C   sH   | j }| j}| }t||jdd t|j|dd t|j|dd d S )NgHz>r   r   r   )r   r#   r   r   Z
predict_mur   r   )r   r   r   r   r    r    r!   test_predict  s    z%TestGLMLogitConstrained2.test_predictc                 C   s4   | j  }td|jv  t| j j}|dks0J d S )Nr_   zx1 - x3 = 0.0)r#   r`   r   ra   strrr   )r   rc   Z	lc_stringr    r    r!   rd     s    
z%TestGLMLogitConstrained2.test_summaryc                 C   s\   dd l }| & |dt | j }W d    n1 s<0    Y  td|jd v  d S )Nr   r   r_   )r3   r   r   r   r#   re   r   ra   )r   r3   rc   r    r    r!   rf     s
    
(z&TestGLMLogitConstrained2.test_summary2c                 C   sJ   | j }ddlm} || jj| j}t|j|jdd t|j|jdd d S )Nr   )fit_constrained_wrapr   r   )r   r   r   r#   Zmodelr   r   r   )r   r   r   Zres_wrapr    r    r!   test_fit_constrained_wrap  s
    z2TestGLMLogitConstrained2.test_fit_constrained_wrapN)r7   r8   r9   rg   r^   r   r.   rh   ri   rd   rf   r   r    r    r    r!   r     s   



r   c                   @   s   e Zd Zedd ZdS )TestGLMLogitConstrained2HCc                 C   s   t d | _tj| _ttjtjt	
 d}d}ddi}d}|j|||dd| _| jj\}}t|||d||dd	| _||f| _d S )
Nr   ZHC0Zscaling_factorgB!?r   r   )cov_typecov_kwdsr   )r   r   r   ro   )r   r   r   Zresults_constraint2_robustr   r   r   rQ   r   r   r   r   r#   rr   r   r   )rY   r   r   r   r\   r   r   r    r    r!   r^     s"    


z&TestGLMLogitConstrained2HC.setup_classNrt   r    r    r    r!   r     s   r   c            
      C   s   d} t j| ttd jd}|  d}dd l}||j|}|j	|j
 }}|j||ddid d	}t j|td
}|  ||jd}	|j|	j	|	j
ddid d S )Nz*deaths ~ C(agecat) + C(smokes) : C(agecat)r   rm   z:C(smokes)[T.1]:C(agecat)[3] = C(smokes)[T.1]:C(agec`at)[4]r   rD   rB   ro   r>   r?   rA   rn   )r	   rL   r@   rq   ry   rS   rT   rN   rU   rV   rW   r   )
Zformula2r[   rr   rS   r]   r   r   Z	formula1aZmod1aZlc_1ar    r    r!   junk  s&    r   ):__doc__ior   Znumpyr   Znumpy.testingr   r   r   ZpandasZpdrS   r.   Zstatsmodelsr   r   r   Z#statsmodels.discrete.discrete_modelr	   r   r   r   r   Zstatsmodels.tools.toolsr   rJ   r   r   r   Zspectorloadr   ZasarrayrQ   r   r2   ZssZread_csvr@   ZastypeintrP   objectr   r:   rj   ru   rx   rz   r|   r   r   r   r   r   r   r   r   r    r    r    r!   <module>   sJ   

H.!6;