U
    Kvf%                     @   s$  d dl Zd dlZd dlZd dlmZmZmZ d dl	m
Z
 d dlmZ d dlmZ ejdddd	gddd
dgddddgddddgddddgddddgdddd	gddddgddddgddddgdddd gdd!d"d#gdd$d%d ggd&d'd(d)gd*Zd+d, Zd-d. Zejjd/d0 Zd1d2 Zd3d4 Zd5d6 ZdS )7    N)assert_almost_equalassert_raisesassert_allclose)MANOVA)MultivariateTestResults)add_constantzMinas Graesg%C @g(\ @gHzG?gd;O @goʡ?gQ @g+?g- @gX9v @gV- @g      @gS㥛?gQ@g^I+@gL7A`?zMatto Grossog\(\ @g;On @gS㥛 @gNbX9 @g}?5^I?gS @gx @z
Santa CruzgbX9 @g r?g @g+ @g+?gE @g+ @LocBasalOccMax)columnsc                  C   sJ  t jdtd} |  }t|d d jd ddd t|d d jd	 d
dd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd ddd t|d d jd  d!dd t|d d jd" d#d$d t|d d jd% d&dd t|d d jd' d(d$d t|d d jd) d*d$d t|d d jd+ d,d$d t|d d jd- d.d$d d S )/NBasal + Occ + Max ~ Locdatar   statWilks' lambdaValue%R>?   decimalPillai's tracer   !?Hotelling-Lawley tracer   /z?Roy's greatest rootr   3^a?r   F Valuep=
ף?   r   r"   Q?r   r"         ?r   r"   Q?r   Num DF      r   r,   r   r,   r   r,   r   Den DF   r   r3      r   r3   iq."@   r   r3   	   r   Pr > FvOjM?r   r=   48E?r   r=   n?r   r=   #J{/L?)r   from_formulaXmv_testr   loc)modr rK   N/tmp/pip-unpacked-wheel-2v6byqio/statsmodels/multivariate/tests/test_manova.pytest_manova_sas_example   s                        rM   c                  C   s  t tjtdg dtd} tdddg }t|| }td}d|d	< td
}d |d< |d< d|fd|fg}||}t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d ddd t	|d d j
d  d!dd t	|d d j
d" d#d$d t	|d d j
d% d#d$d t	|d d j
d& d#d$d t	|d d j
d' d$d$d t	|d d j
d( d)d$d t	|d d j
d* d+d$d t	|d d j
d, d-d.d t	|d d j
d/ d0d$d t	|d d j
d1 d2d.d t	|d d j
d3 d4d.d t	|d d j
d5 d6d.d t	|d d j
d7 d8d.d d S )9Nr   TZ
drop_firstZdtyper	   r
   r   )   r.   rO   )r   r   )r$   r.   )r   rO   )rO   r$   Z	Interceptr   r   r   r   r   r   r   r   r   r   r    r!   r#   r$   r%   r&   r'   r(   r)   r*   r+   r-   r.   r/   r0   r1   r2   r4   r5   r6   r7   r8   r9   r:   r;   r<   r>   r?   r@   rA   rB   rC   rD   )r   pdget_dummiesrF   floatr   npzerosrG   r   rH   )exogendogrI   Z	interceptrH   Z
hypothesesrJ   rK   rK   rL   test_manova_no_formulaO   s    



                    rW   c                  C   sL   t tjtdg dtd} tdddg }t|| }| }t|tsHt	d S )Nr   TrN   r	   r
   r   )
r   rP   rQ   rF   rR   r   rG   
isinstancer   AssertionError)rU   rV   rI   rJ   rK   rK   rL   $test_manova_no_formula_no_hypothesis   s    
rZ   c                  C   s   t jdtd} dtdddggd fg}| | dtddggd fg}tt| j| dtdddggtdgdgdggfg}| | dtdddggtdgdggfg}tt| j| d S )Nr   r   testrO   )r   rE   rF   rS   arrayrG   r   
ValueError)rI   Z
hypothesisrK   rK   rL   !test_manova_test_input_validation   s    
,
(r^   c                   C   s   t ttjdt d S )NzBasal ~ Loc)r   r]   r   rE   rF   rK   rK   rK   rL   test_endog_1D_array   s    r_   c                  C   s  d} dddg|  }t j| ddddg  }t j| dd	d
dg  }t j| ddddg  }tt||||d}tjd|d}|	 }|dddg 
 }|dddg | |dddg< tjd|d}|j	dd}	|jd d t}
|	jd d t}t|
|dd d S )N   r	   r
   r   r.   r   g      ?rO   g      ?r(   g333333?g333333?)r   r	   r
   r   z#Basal + Occ + Max ~ C(Loc, Helmert)r   T)Zskip_intercept_testzC(Loc, Helmert)r   g|=)Zrtol)rS   randomZrandnZravelrP   	DataFramedictr   rE   rG   ZmeanresultsZto_numpyrR   r   )ZngrH   y1y2Zy3ZdtarI   Zres1ZmeansZres2Zstat1Zstat2rK   rK   rL   test_manova_demeaned   s    rg   )ZnumpyrS   ZpandasrP   ZpytestZnumpy.testingr   r   r   Zstatsmodels.multivariate.manovar   Z)statsmodels.multivariate.multivariate_olsr   Zstatsmodels.toolsr   rb   rF   rM   rW   markZsmokerZ   r^   r_   rg   rK   rK   rK   rL   <module>   s:   












0;

