U
    md                     @   s  d dl mZ d dlZd dlZd dlZd dlZd dlZd dl	m
Z
 d dlZd dlmZmZmZmZ d dlmZmZmZmZ d dlmZ d dlm  mZ ddlmZ G d	d
 d
Zdd Z G dd dZ!G dd dZ"G dd de"Z#dd Z$ej%&ej%'e(Z)ej%*e)dZ+e,e+Z-dd e-D Z-ej./de-ej./dddgej./dddgdd Z0dd  Z1d!d" Z2ej.j3d#d$ Z4d%d& Z5ej.j3d'd( Z6d)d* Z7G d+d, d,Z8d-d. Z9G d/d0 d0Z:d1d2 Z;d3d4 Z<dS )5    )PLATFORM_OSXN)sparse)MixedLMMixedLMParams_smw_solver_smw_logdet)assert_almost_equalassert_equalassert_allcloseassert_)
_penalties   )lme_r_resultsc                   @   s   e Zd ZdZdd ZdS )	R_ResultsaY  
    A class for holding various results obtained from fitting one data
    set using lmer in R.

    Parameters
    ----------
    meth : str
        Either "ml" or "reml".
    irfs : str
        Either "irf", for independent random effects, or "drf" for
        dependent random effects.
    ds_ix : int
        The number of the data set
    c              	   C   s  d|||f }t td| | _t td| | _t td| | _t td| | _t td| | _ttd| rt td| | _t td| | _	t
| j	| _	tjtjt}tj|d	}tj|d
| }t|dd&}t|}	t|	}
dd |	D }W 5 Q R X t
|}|d d |
df | _|d d |
df | _dd t|
D }|d d |f | _dd t|
D }|d d |f | _d S )Nz	_%s_%s_%dcoefZvcovcov_rescaleloglikeZ
ranef_meanranef_condvarresultszlme%02d.csvzutf-8)encodingc                 S   s   g | ]}d d |D qS )c                 S   s   g | ]}t |qS  )float.0xr   r   ^/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/statsmodels/regression/tests/test_lme.py
<listcomp>@   s     z1R_Results.__init__.<locals>.<listcomp>.<listcomp>r   )r   liner   r   r   r   @   s     z&R_Results.__init__.<locals>.<listcomp>endoggroupsc                 S   s   g | ]\}}| d r|qS )exog_fe
startswithr   ir   r   r   r   r   F   s     
 c                 S   s   g | ]\}}| d r|qS exog_rer"   r$   r   r   r   r   H   s     
 )getattrr   r   vcov_rcov_re_rscale_rr   hasattrranef_postmeanr   npZ
atleast_2dospathdirnameabspath__file__joinopencsvreadernextasarrayindexr   r    	enumerater!   r'   )selfmethirfsds_ixZbnamecur_dirrdirfnamefidrdrheaderdataiir   r   r   __init__)   s4    

zR_Results.__init__N)__name__
__module____qualname____doc__rH   r   r   r   r   r      s   r   c                    s    fdd}|S )Nc                    s*   t j| jjj d}j|d S )N)has_fe
profile_fe)r   Zfrom_packedk_fek_reuse_sqrtr   )r   paramsrM   modelrO   r   r   fP   s        zloglike_function.<locals>.fr   )rU   rO   rM   rV   r   rT   r   loglike_functionL   s    rW   c                   @   s   e Zd Zejjejdddgejdddgejdddgdd Zdd	 Zd
d Z	ejjejj
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zejjdd Zdd Zejjd d! Zd"S )#TestMixedLMrR   FTremlrO   c           $   	   C   s.  d}d}d}d}t jd t jj|| |fd}t jj|| |fd}	d|	d d df< t jj|| dfd}
t jj||fd}|d d d	f  d9  < t |t |df}t jj|dfd}t |t |df}|d d d	f  d9  < ||	 d}||
 d}t jj|| d}|d| | | }t t|t |}i i d
}t|D ]D}t ||k}|
|ddf |d |< |
|ddf |d |< qZt	j
tdd t||||	||d}W 5 Q R X |j|d}t||| d}ztdD ]}t jj|d}t jj||fd}t |j|}t jjddd }tj|||d}|j| |d}|j||d }t||}t||dd q|dkr|dkr||j\} }!|!rt	d | d	9 } |jjddd}t|ddd}"t||"}#t| |#dd W n* tk
r(   tr"t	d n Y nX d S )N            i  sizer   r   abrb   rc   Using deprecated variancematch)exog_vcrR   rY   )rO   rM   )r   vcomprM   rR   rN   MbP?rtolFzhessian should not be singularT)rR   rM   z4fails on OSX due to unresolved numerical differences)r.   randomseednormalkrononessumrangeflatnonzeropytestwarnsUserWarningr   fitrW   dotTr   from_components
get_packedZscorendZapprox_fprimer
   Zhessianparams_objectfailZapprox_hessAssertionErrorr   Zxfail)$r<   rR   rY   rO   n_grpZgrpsizerP   rQ   r!   r'   rg   slopesZ	slopes_vcZ	re_valuesZ	vc_valueserrr   r    vcr%   ixrU   rsltr   kr	fe_paramsr   ri   rS   Z
params_vecgrZngrZhessZsingZ	loglike_hZnhessr   r   r   test_compare_numdiff[   s    
     
   z TestMixedLM.test_compare_numdiffc                 C   s   t jd t jjdd}t t ddddg}t t jjdddddg}|d| t jjdd }t||| }t|||t 	d }t
|j|jdd d S )	N  ,     r^   d   r   r      decimal)r.   rn   ro   rp   rq   arangers   r   ry   rr   r   rS   )r<   exogr    g_errorsr   mdf1mdf2r   r   r   test_default_re   s    zTestMixedLM.test_default_rec                 C   s   t jd t jjdd}t t ddddg}t t jjdddddg}|d| t jjdd }t|||}|jdd}t	t
|d	d d S )
Nr   r   r^   r   r   r   T)Zfull_outputhist)r.   rn   ro   rp   rq   r   rs   r   ry   r	   r,   )r<   r   r    r   r   modr   r   r   r   test_history   s    zTestMixedLM.test_historyc              	   C   s  t jd d}d}d}t jj|| |fd}t || df}t t |t |}t jj|| d}t jj|| d}d}	t t jjddt |}
|	|
|d d df  7 }	t jj|d}|	t |t || 7 }	t jj|d}|	t |t || 7 }	|	t jj|| d7 }	|d|	 }i i d}t|D ]L}t 	||k}|| d d d f |d	 |< || d d d f |d
 |< q:t
jtdd t|||||d }W 5 Q R X |jddddddd |jd
dddddd d S )NiV&  r]   r\   r   r^   r   r   ra   rb   rc   rd   re   )r    r'   rg   re)ZvtypeZdist_lowZnum_lowZ	dist_highZnum_highr         ?)r.   rn   ro   rp   rr   rq   r   rs   rt   ru   rv   rw   rx   r   ry   Z
profile_re)r<   rP   gsizer   r   r'   r    ZvcaZvcberrorsr   rcr   r   krG   r   r   r   r   test_profile_inference   sH    
    z"TestMixedLM.test_profile_inferencec              	   C   s  t jd t jjdd}t jjdd}t t dt d}t jjdd}|d d df  d	9  < t |t d
| }|dt jjdd }|d| }tdd	d}t d|_	t 
d	|_t d|_t||||d}|j|d}	i i d}
t|jD ]D\}}|j| }||ddf |
d |< ||dd	f |
d |< qtjtdd t||||
d}W 5 Q R X | }|  t|	j	|j	dd tt |	j|jddd t|	jdddg |jddd d S )N  )  r   r^   )r   r]   r   r   )r   r]   r   r]   r   r   r   r   r&   )freera   rb   rc   rd   re   )rg   -C6?)atol{Gz?)r   rm   r\   )r.   rn   ro   rp   rq   r   rr   rs   r   r   eyer   zerosri   r   ry   r;   Zgroup_labelsZrow_indicesrv   rw   rx   summaryr
   diagbse)r<   r   r'   r    r   r   r   r   model1result1rg   r   groupr   model2result2r   r   r   test_vcomp_1   sJ    


      zTestMixedLM.test_vcomp_1c                 C   sF  t jd d}t jj|dfd}t t |d t d}d}t jj|dfd}t jj|d dfd}t |t d| }||d7 }t t |d	 t d	}|t dt jj|d	 d t d	7 }t t |d t d}|t dt jj|d d t d7 }|t jj|d7 }|d| }	tj	t
|d
}
|	|
d< ||
d< |d d df |
d< |d d df |
d< |d d df |
d< |d d df |
d< ||
d< ||
d< ddd}tjd|d||
d}| }t|jdddgdd t|jddgdd ggd!d t|jd"d#gd!d t|jjdd$ d%d&d'gd!d d S )(Na  i@  r]   r^      r   )r   r   r   r   r:   yr    x1x2z1z2v1Zv2z	0 + C(v1)z	0 + C(v2)ra   zy ~ x1 + x2z0+z1+z2r    
re_formula
vc_formularF   gE>'?gh㈵?g|'f?r   rl   gCl?g㥛 ?gV-?rk   gjt@g`"@r\   gFx$?gD;)?g?)r.   rn   ro   rp   rq   r   rr   rs   pd	DataFramert   r   from_formulary   r
   r   r   ri   r   iloc)r<   nr   r    r   r'   r   Z
subgroups1Z
subgroups2r   dfvcfr   r   r   r   r   test_vcomp_2  sb    ((
      zTestMixedLM.test_vcomp_2c           	      C   s   t jd t jjdd}t t dt d}t jjdd}t |t d| }|t jjdd }ddi}t|||d}t	j
d	d
||d}| }|  t|jjdd t jd dd t|jjdd t jd dd d S )Nr   r   r^   r   r   rb   z0 + x1)r   r   r    zy ~ 1r    )r    r   rF   r   )gSvA]g{?gб?gP4`_rk   rl   )ghg 8L?g!Va̿g&R?o*)r.   rn   ro   rp   rq   r   rr   r   r   r   r   ry   r   r
   residr   r_fittedvalues)	r<   r   r    r   r   Zvc_fmlr   rU   resultr   r   r   test_vcomp_3Q  s4       zTestMixedLM.test_vcomp_3c           
      C   s   t jt jt}t j|d}t j|d}t|}ddi}tj	ddd||d}|
 }tj	ddd|d	|d
}|
 }	t|j|	j t|j|	j d S )Nr   
pastes.csvcask0 + caskstrength ~ 1batch1r   T)r    r   r   Z
use_sparserF   )r/   r0   r1   r2   r3   r4   r   read_csvr   r   ry   r
   rS   r   )
r<   r@   rA   rB   rF   r   rU   r   r   r   r   r   r   test_sparsek  s0    
zTestMixedLM.test_sparsec                 C   sP  t jt jt}t j|d}t j|d}t|}tj	dd|d}|
 }t|jtjd dd t|jd	d
 tjd dd t|jddd t|jddd t||jddd t|}tj	dd|d}|j
dd}t|jtjd dd t|jd	d
 tjd dd t|jddd t|jddd t||jddd d S )Nr   
dietox.csvWeight ~ TimePigr    rF   )gMqr/@gD @h㈵>rl   r   r]   )gKz7?gd"?g$0ܻ&@g-!l2D@ǧɢFrh   )g73pr/@g"J @)gǎ?g\?g0|&@gC@g/Ţ)r/   r0   r1   r2   r3   r4   r   r   r   r   ry   r
   r   r.   r   r   r   r   r   r   r<   r@   rA   rB   rF   rU   r   r   r   r   test_dietox  sV    
    
  
    
  zTestMixedLM.test_dietoxc                 C   sr  t jt jt}t j|d}t j|d}t|}tj	ddd|d}|j
dd}t|jtjd	 d
d t|jdd tjd dd t|jddd t|jj tjd dd t||jdd
d t|}tj	ddd|d}|j
ddd}t|jtjd d
d t|jdd tjd dd t|jddd t|jj tjd dd t||jdd
d d S )Nr   r   r   r   z1 + Time)r    r   rF   Zcg)method)g|a2U0z/@g@r   rl   r   r]   )gUx_?gcoo?rk   gJY&@)g
zS~3@矊%?r   g=Ƃ ?皙?gMbRF)r   rY   )g`-z/@g(yu@)gS~;{?gjt1K?g?V&@)g@C03@wo?r   gl`q8?r   gK7O)r/   r0   r1   r2   r3   r4   r   r   r   r   ry   r
   r   r.   r   r   r   r   valuesZravelr   r   r   r   r   r   test_dietox_slopes  sr    
       

  
     

  zTestMixedLM.test_dietox_slopesc           
      C   s  t jt jt}t j|d}t j|d}t|}ddi}tj	ddd||d}|
 }t|jjd	 d
dd t|jjd	 ddd t|jjd ddd t|jddd t|jddd t|jtj t|jtj tjd }t|jd	d |dd tjd }	t|jd	d |	dd tj	ddd||d}|j
dd}t|jjd	 d
dd t|jjd	 ddd t|jjd ddd t|jddd t|jddd t|jddd t|jddd d S )Nr   r   r   r   r   r   r   r   r   g;pΈN@rk   rl   ga2U0*?)r   r   gPn?gV-?g(\^r   )gVQ?g7BZg4Kbg0_^?gDr[   )&NwPO@r   9}N@r   g7ُiO@r   Frh   g%C?gv/?g5{?g+^go@grp@)r/   r0   r1   r2   r3   r4   r   r   r   r   ry   r
   r   r   r   r   r   llfr	   Zaicr.   nanZbicr   r   r   )
r<   r@   rA   rB   rF   r   rU   r   r   ry   r   r   r   test_pastes_vcomp  sN    	


zTestMixedLM.test_pastes_vcompc              	   C   s  t jd d}t jj|dfd}d|d d df< g }t t |d t d}d}t jj|dfd}t jj|d dfd}t |t d| }||d7 }t jj|dfd}t jj|d dfd}|d d dd f  d9  < t |t d| }||d7 }|t jj|d7 }|d| }i i d	}	tt	t
|d D ]H\}
}t ||k}||ddf |	d
 |< ||dd f |	d |< qHtjtdd t|||||	d}W 5 Q R X | }tj|d d dd f dgd}||d< |d d df |d< |d d df |d< |d d df |d< |d d df |d< |d d df |d< |d d df |d< ddd	}tjd|d||d}| }t|j|jdd t|j|jdd t|j|jdd t|j|jdd t|j|jdd d S ) Nr   i   r]   r^   r   r   r   r   ra   rb   rc   rd   re   )r'   rg   r   columnsr   Zre1Zre2Zvc1Zvc2Zvc3r\   Zvc4z0 + vc1 + vc2z0 + vc3 + vc4y ~ x1z0 + re1 + re2r   g:0yE>rl   )r.   rn   ro   rp   rq   r   rr   rs   r;   rt   intru   rv   rw   rx   r   ry   r   r   r   r
   r   r   ri   rS   r   )r<   r   r   Zex_vcr    r   r'   r   r   rg   r   r   r   r   r   r   r   r   r   r   r   r   test_vcomp_formula^  sd    

 
zTestMixedLM.test_vcomp_formulac              	   C   s  t jd t jjdd}t jjdd}t t ddddg}|t t jjdddddg }|d| t jjdd }t||||}t|j	j
ddd	d
gk t|j	jdgk t|j	jdgk | }td|i}t|jd D ]}	|d d |	f |d|	 < q||d< d}
d}tj|
|||d}t|j	j
ddddgk t|j	jdgk t|j	jdgk | }t|j|j ||d< tj|
||dd}t|j	j
ddddgk t|j	jdgk t|j	jdgk |j|jd}t|j|jdd t jt|t jd}t||||}t  td | }W 5 Q R X ddlm} ||
|dd }t|j	jdgk t|j	jd!gk t  td | }W 5 Q R X t|j|j d S )"Nj	  r   r^   r   r   r   r   r   Zx3Zx4Zx_re1z	x_re1 Varr   exog%dr'   )endog ~ 0 + exog0 + exog1 + exog2 + exog3z0 + exog_rer   r    exog0exog1exog2exog3exog_re Varr    Zstart_paramsr   rl   Zdtypeignorer   )mixedlmr    z
groups Var)r.   rn   ro   rp   rq   r   rs   r   r   rF   xnamesexog_re_namesexog_re_names_fullry   r   r   rt   shaper   r   rS   r
   rr   lenZfloat64warningscatch_warningssimplefilterZstatsmodels.formula.apir   )r<   r   r'   r    r   r   mod1Zrslt1r   r   fmlre_fmlmod2Zrslt2Zmod3Zrslt3mod4Zrslt4r   Zmod5Zrslt5r   r   r   test_formulas  s`        



zTestMixedLM.test_formulasc              	   C   sH  t jd t jjdd}t t dt d}|d d df |d d df  }|t t jjddt d t jjdd }t|||}|jd	d
}|	  t|||}|jdt d d
}|	  t
 }|j|dd}	|		  t
 }t  td |j|dd}
W 5 Q R X |
	  t
d}|j|d	d}|	  d S )Ni}  )r   r[   r^   r   r   r   r]   r   g      ?)alpha
   r[   g        )r   r
  r   g      $@g333333?)r.   rn   ro   rp   rq   r   rr   r   fit_regularizedr   	penaltiesZL2r  r  r  ZPseudoHuber)r<   r   r    Zexpected_endogr   mdr   r   penZmdf3Zmdf4Zmdf5r   r   r   test_regularized  s4     


zTestMixedLM.test_regularizedN)rI   rJ   rK   rv   markslowparametrizer   r   r   Zsmoker   r   r   r   r   r   r   r   r   r	  r  r   r   r   r   rX   X   s,   T!&?>ER
6;rX   c                   @   s0   e Zd Zedd Zdd Zdd Zdd Zd	S )
TestMixedLMSummaryc           
   
   C   s   t ddgd}t dgd}ddddddddd	d
g
}dddddddd
ddg
}ddddddddd
d	g
}t|||||d}|d j}|dddg j}|d j}	t|||	d | _d S Nr   r   r[   r     r\   r]      	   r   r   )r   pidx0r   r   r   r  r   r   r  r   )r.   repeatr   r   r   r   ry   res
clsr  r  r   r   r   r   r   r   r    r   r   r   setup_class  s    

zTestMixedLMSummary.setup_classc                 C   s2   | j  }ddddg}|jd jj}t|| d S )Nconstr   r   	Group Varr   r  r   Ztablesr:   r   r	   r<   ZsummZdesiredactualr   r   r   test_summary
  s    
zTestMixedLMSummary.test_summaryc                 C   s<   | j jdddgd}ddddg}|jd jj}t|| d S )NConstantZAgeZWeight)Zxname_fer!  r   r"  r#  r   r   r   test_summary_xname_fe  s    z(TestMixedLMSummary.test_summary_xname_fec                 C   s8   | j jdgd}ddddg}|jd jj}t|| d S )NzRandom Effects)Zxname_rer   r   r   r   r"  r#  r   r   r   test_summary_xname_re  s    z(TestMixedLMSummary.test_summary_xname_reN)rI   rJ   rK   classmethodr  r%  r'  r(  r   r   r   r   r    s
   
r  c                   @   s   e Zd Zedd ZdS )TestMixedLMSummaryRegularizedc           
   
   C   s   t ddgd}t dgd}ddddddddd	d
g
}dddddddd
ddg
}ddddddddd
d	g
}t|||||d}|d j}|dddg j}|d j}	t|||	d | _d S r  )r.   r  r   r   r   r   r  r  r  r   r   r   r  )  s    

z)TestMixedLMSummaryRegularized.setup_classN)rI   rJ   rK   r)  r  r   r   r   r   r*  &  s   r*  c              	   C   s  |r|dk rd S |rdnd}| r$dnd}t |||}t|j|j|j|j}|stt|j	dk rt
  t
d |jd| d	}W 5 Q R X n|jd| d	}n|jjd
 }|jjd
 }	t||	d}
t||
_t|	|
_tg |
_tt|j	dk r2t
   t
d |j| d|
d}W 5 Q R X n|j| d|
d}t|j|jdd t|j|j	dd t|j|jdd |j}t|j| d|d|f dd t|j|jd dd |st|jd |j dd t|j!d |j"dd d S )Nr  irfZdrfrY   mlr   r   Hz>)gtolrY   r   r   )rY   r.  r   r   r   r\   r]   )#r   r   r   r!   r    r'   r.   anyr   r*   r  r  r  ry   r   r   rr   r   r   r   arrayri   r   r   r   r+   rP   r)   
cov_paramsr   r   random_effectsr-   random_effects_covr   )rY   r+  r?   r>   r=   r   r  mdfrP   rQ   r   r   r   r   do1<  sX    



      r5  r   c                 C   s$   g | ]}| d r|dr|qS )Zlmez.csv)r#   endswithr   r   r   r   r   x  s     
 
 r   rB   rY   FTr+  c                 C   s    t | dd }t||| d S )Nr\   r[   )r   r5  )rB   rY   r+  r?   r   r   r   test_r|  s    r7  c                  C   s^  t jd t jjdd} t jjdd}t t ddddg}|t t jjdddddg }| d| t jjdd }td|i}t	| j
d D ]}| d d |f |d| < q||d	< d
}d}tj||||d}	|	 }
|
  ddddg}dd	g}dddg}t|	jj|k t|	jj|k t|	jj|k |
j}t|j || k |
j}t|j || k |
j}t|j || k |
 }t|j || k t|j || k |
j}t|j |k |
j}t|j |k |
j}t|j |k t|j |k |
j}t|j |k t|j |k |
j}t|j |k d S )Nr   r   r^   r   r   r   r   r   r'   r   z	~ exog_rer   r   r   r   r   Groupr!  zGroup x exog_re Covr   ) r.   rn   ro   rp   rq   r   rs   r   r   rt   r   r   r   ry   r   r   rF   r   r   r   rS   r:   tolistr   tvaluesr1  r   r   bse_fer   Zcov_re_unscaledbse_re)r   r'   r    r   r   r   r   r  r  r  r   r   Zre_namesZre_names_fullrS   r   r:  r1  Zfer;  r   Zcov_re_ur<  r   r   r   test_mixed_lm_wrapper  sT     
r=  c                  C   s   t jd d} d}d}d}|t jj| d }t |t |}||t jj| | d  }t t | t |}t | | }t|||d}	|	 }
|
j	}t
t|t t
t|| k t
t|d tj t
t|d d	k t||||d
}	|	 }
|
j	}t
t|t t
t|| k t
t|d tj t
t|d d	k t jj| | dfd}d	|d d df< t dd	|}t t | ||d d d	f< t||||d
}	|	 }
|
j	}t
t|t t
t|| k t
t|d tj t
t|d dk d S )Ni[  r   r  r]   r\   r^   r   r   r   )r'   r    r`   )r.   rn   ro   rp   rq   rr   r   r   ry   r2  r   
isinstancedictr   r   SeriesZlinspace)Zngrpr   ZrsdZgsdZmnZgmnr   r   r   rU   r   r   ZxrZqpr   r   r   test_random_effects  sF    rA  c                  C   s<  t jd t jjdd} t| } dddddd	g| _t t d
t 	d| d< t jjdd}t |t 	d}|d d df |d d df | j
  |d d df | j  | d< | d  |d d df | j t jjdd 7  < t j| jd< t j| jd< t j| jd< t j| jd< t j| jd< t j| jd< d}d}ddd}dD ]}dD ]}i }|  }dddg}	|rz||d< |	d |r||d< |	dd	g ||	  }t  td  tj|fd|d!|}
|
 }tj|fd| d"d#|}| }t|j|j t|j|j tt|j|j W 5 Q R X qHq@d S )$Ni[  )r   r  r^   r   gr   r   c1c22   r]   )rE  r   )r]   r   r   r   r\   r   )r   r   )r]   rB  )r\   r   )r   r   )r[   rC  )r  rD  r   z1 + z1z0 + c1z0 + c2ra   )FTr   r   r   r   Zdrop)r    rF   missing) r.   rn   ro   rp   r   r   r   rq   r   rr   r   rC  rD  NaNloccopyappendextendZdropnar  r  r  r   r   ry   r
   rS   r   r	   r   r   Znobs)r   r   r  r   r   Z
include_reZ
include_vckwargsZdxvar   r   r   r   r   r   r   test_handle_missing  sl    
@0




   rN  c                  C   s  ddl m}  dddddddddddddddg}ddddddd	d
dddddddg}tddddddddddddddd g}|||d!}t|}td"d#d$g}tjd%||d& d'}|j|d(}td)d*d+g}tjd,||d& d'}	|	j|d(}
| ||
gd-d.}d/}t	t
|| d S )0Nr   )summary_colr   r]   r\   r   r[   r  r  r   r  r                 goʡ?gd;OgZd;@gZd;O@gjt@gl@grh@gGz@g"~@g(\O$@g$@gDl)1@g-F1@g`"-@g-,@)YXIDSg@{Fg~eJ?gͣT?Y ~ XrW  r   r   g&ʓ@g.?g<H.*?zX ~ YT)starsa  
=============================
              Y         X    
-----------------------------
Group Var 0.1955    1.3854   
          (0.6032)  (2.7377) 
Intercept -1.2672   3.4842*  
          (1.6546)  (1.8882) 
X         1.1618***          
          (0.1959)           
Y                   0.5529***
                    (0.2080) 
=============================
Standard errors in
parentheses.
* p<.1, ** p<.05, ***p<.01)Zstatsmodels.iolib.summary2rO  r.   r0  r   r   r   r   ry   r	   str)rO  Zidsr   r   dZsp1r  Zresults1Zsp2r  Zresults2outsr   r   r   test_summary_col   s>    ""             
r^  c                     s  t jd d} d}g g g g g g g g g f	\	}}}}}}}}	}
t| D ]f}t jj|dfd}|d d df dt jj|d  }t jj|dfd}t jjdd}|d  d9  < |d	  d	9  < |t ||j7 }|| t t j	d
 t 
|d t}t jjdd}||| 7 }|| |	| t t 
|d t j	d
 t}t jjdd}||| 7 }|| |
| || || || |d| g|  qDt |}t |}t |}t |}t |}t |}t||d d df |d d d	f |d d df |d d d	f |||d}t |}t |	}	t |
}
t j|	|
fd	d}tjdddddd|d}| }|j  fddt| D }t |}t |d d df |d d df d
 dkstt |d d d	f |d d d	f d
 dkst fddt| D }t |}tdD ]8}t |d d |f |d d |f d
 dks:tq:|j}| D ](} | j}|| j||fkstqd S )Ni  i  r  r]   r^   r   r   r\   r   )r   r   g%d)r   r  r   z0r   v0r   rB  )Zaxiszy ~ x0 + x1z~0 + z0 + z1z~0+C(v0)z0+C(v1))ra  r   rB  )r   r   r    rF   c                    s    g | ]} d |  dd qS )r_  r   r]   r   r   r   refr   r   r     s     z/test_random_effects_getters.<locals>.<listcomp>g?c                    s    g | ]} d |  dd qS )r_  r]   r  r   rb  rc  r   r   r     s     r   )r.   rn   ro   rt   rp   rz   ZflatrJ  rq   r   rr   Zastyper   Zconcatenater   r   r9   r   r   ry   r2  Zcorrcoefr   r3  keysr_   r   )ngmr   r   zra  r   rB  rc   Zc0rC  r%   ZxxyyzzZbbZvv0Zcc0Zvv1Zcc1r   ccrU   r   Zb0Zcf0r   Zrefcpr   rc  r   test_random_effects_gettersB  s    ("
"

"













	
22
6
rm  c              
   C   s6  || }t jj| |fd}t |j|}t ||f}t jj||fd|d|d|f< t jj|d}t d| |||||f< t j	|d|d|f }	d}t jj| d}
t j
|t | |  t |t ||j |
}t||||	|}||
}t|| t|t|t||	|}||
}t|| d S )Nr^   r   r   r   )r.   rn   rp   rz   r{   r   uniformr   linalginvZsolver   r   r
   r   Z
csr_matrix)rl  qrr]  r[  AAtABdiQir   y2rV   y1r   r   r   check_smw_solver  s&    "0
rz  c                   @   sh   e Zd Zedd Zejdddgejdddgejd	d
dgejdddgdd ZdS )TestSMWSolverc                 C   s   t jd d S N   r.   rn   ro   r  r   r   r   r    s    zTestSMWSolver.setup_classrl  r[   r  rq  r   r   rr  r]   r\   r]  r   r   c                 C   s   t |||| d S N)rz  r<   rl  rq  rr  r]  r   r   r   test_smw_solver  s    zTestSMWSolver.test_smw_solverN)	rI   rJ   rK   r)  r  rv   r  r  r  r   r   r   r   r{    s   
r{  c              
   C   s&  || }t jj| |fd}t |j|}t ||f}t jj||fd}t |j||d|d|f< t jj|d}	t d|	 |||||f< t j	|d|d|f }
d}t j
|t | |  t |t ||j \}}t j
|\}}t||||
|	|}trdnd}t|||d d S )Nr^   r   r   r   gư>r-  rl   )r.   rn   rp   rz   r{   r   rn  r   ro  rp  Zslogdetr   r   r   r
   )rl  rq  rr  r]  r[  rs  rt  ru  crv  rw  _Zd2ZbdZd1rm   r   r   r   check_smw_logdet  s    2r  c                   @   sh   e Zd Zedd Zejdddgejdddgejd	d
dgejdddgdd ZdS )TestSMWLogdetc                 C   s   t jd d S r|  r~  r  r   r   r   r    s    zTestSMWLogdet.setup_classrl  r[   r  rq  r   r   rr  r]   r\   r]  r   r   c                 C   s   t |||| d S r  )r  r  r   r   r   test_smw_logdet  s    zTestSMWLogdet.test_smw_logdetN)	rI   rJ   rK   r)  r  rv   r  r  r  r   r   r   r   r    s   
r  c               	   C   s   t jd d} t j| d}tj|ddgd}tjdd |jD |jd	|d
< t	t
8}tjd||d
 d}| }|  |std W 5 Q R X d S )Ni_  r   r]   rU  rV  r   c                 S   s   g | ]}|d  qS )r\   r   )r   r%   r   r   r   r     s     z!test_singular.<locals>.<listcomp>r   classrX  r   zwarning expected)r.   rn   ro   Zrandnr   r   r@  r:   rv   rw   Warningr   r   ry   r   r   )r   rF   r   Zwrnr  r4  r   r   r   test_singular  s     r  c               	   C   s   t jd d} d}t jd }t ddgddgg}t jd }d	}t jj| dfd
}t jj| dfd
}t jj| dfd
}t jj| dfd
}	t t j|td| | }
t 	||}t jj|dfd
}t 	|t j
|j}t |d t jj|dfd
 }t |d t jj|dfd
 }|||
d d f | d }|||
d d f | d7 }|||
d d f |	 d7 }|t |t jj| d
 7 }t||d d df |d d df |d d df |d d df |
d}|d d df |d< |d d df |d< |	d d df |d< |	d d df |d< ddd}tjddd||d}dd l}|jjj}|j|||d}|jddd}|t|d   |  < |||d }|d}t 	||| d }||d |d d  7 }||d |	d d  7 }||7 }tt || |ddd d S ) N   r   r  )r   r   r   r]   )g      ?g      @g      ?r^   r   r   )r   r   r   r`  r   grpr   Zz3Zz4Zz5z0 + z2 + z3z0 + z4 + z5ra   zy ~ 0 + x1 + x2r  z0 + z0 + z1r   )r   r   ri   TFrj   r   r   )rm   r   )r.   rn   ro   r   r9   rp   r  r   r   rz   ro  Zcholeskyr{   sqrtrs   r   r   r   r   statsmodelsZ
regressionZmixed_linear_modelr   r|   r}   r   get_distributionZrvsZmeanr
   var)r   Zn_groupsr   r   ri   r   r!   r'   Zexog_vcaZexog_vcbr    Zeyuu1u2r   r   r   rg  r  mppopadistyrvr   r   r   test_get_distribution   sh    

  " 
 

r  )=Zstatsmodels.compat.platformr   r/   r6   r  numpyr.   Zpandasr   Zscipyr   rv   Z)statsmodels.regression.mixed_linear_modelr   r   r   r   Znumpy.testingr   r	   r
   r   Zstatsmodels.baser   r  Zstatsmodels.tools.numdiffZtoolsZnumdiffr~   r   r   r   rW   rX   r  r*  r5  r0   r1   r2   r3   r@   r4   rA   listdirfnamesr  r  r7  r=  rA  r  rN  r^  rm  rz  r{  r  r  r  r  r   r   r   r   <module>   sX   3     ',9
3-
8"
^