U
    md                     @   sv   d Z ddlZddlZddlmZmZ ddlZddlm	Z	m
Z
 ddlZG dd dZdd Zd	d
 Zdd Zdd ZdS )zTesting helper functions

Warning: current status experimental, mostly copy paste

Warning: these functions will be changed without warning as the need
during refactoring arises.

The first group of functions provide consistency checks

    N)Versionparse)assert_allcloseassert_c                   @   s    e Zd ZdddZdddZdS )	PytestTesterNc                 C   sR   t d}|d kr0|jdd }|d kr0tdtj|| _|jdd | _	d S )N   __file__zUnable to determine path__name__)
sys	_getframef_localsget
ValueErrorospathdirnamepackage_pathpackage_name)selfr   f r   S/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/statsmodels/tools/_testing.py__init__   s    
zPytestTester.__init__Fc                 C   s   zhdd l }t|jtdks t|d kr0ddg}| jg| }tdd|  ||}|rft	
| W n tk
r   tdY nX d S )Nr   z3.0z
--tb=shortz--disable-pytest-warningszRunning pytest  z"pytest>=3 required to run the test)pytestr   __version__r   ImportErrorr   printjoinmainr
   exit)r   
extra_argsr    r   cmdstatusr   r   r   __call__    s    
zPytestTester.__call__)N)NF)r	   
__module____qualname__r   r$   r   r   r   r   r      s   
	r   c                 C   s   | }t t|j}||}t|j|jdd tt |j|j	dd tt |j
|jdd t|j|jdd t| | dd t |j|j	|j|j| f}| j}t||dd tt|d ||d }|  t|j|jd dd d S )N-q=rtolg|=&.>use_tr   )npeyelenparamsZt_testr   ZeffectsqueezesdZbseZtvalueZtvaluespvaluepvaluesZconf_intZcolumn_stackZsummary_framevaluesr   hasattrsummary)resultsresmatttZ	table_resZtable2r   r   r   check_ttest_tvalues0   s$    
 
r;   c                    s   | j tj  fddt D }t|jddd  fddt D }t|jddd dkrvdnd	}t }t||k zt	 }W n t
k
r   Y nX t||k d
S )aN  
    Check that the outputs of `res.wald_test` produces pvalues that
    match res.pvalues.

    Check that the string representations of `res.summary()` and (possibly)
    `res.summary2()` correctly label either the t or z-statistic.

    Parameters
    ----------
    results : Results

    Raises
    ------
    AssertionError
    c                    s(   g | ] }j t | d djqS )T)Zuse_fscalarZ	wald_testr,   r-   r2   .0kk_varsr8   r+   r   r   
<listcomp>`   s   z'check_ftest_pvalues.<locals>.<listcomp>r*   g}:)r)   Zatolc                    s&   g | ]}j t | d djqS )T)r<   r=   r>   )rB   r8   r   r   rC   e   s   FzP>|z|zP>|t|N)r+   r.   r/   ranger   r3   strr6   r   Zsummary2AttributeError)r7   ZpvalsZstring_use_tZsummZsumm2r   rA   r   check_ftest_pvaluesL   s&    
rG   c                 C   s   dd l }ddlm} ddlm} t| d| } t| ||frN|dt	|  | }|j
}t|jj| |jdd t|| dd d S )Nr   
GLMResultsDiscreteResults_resultszNot supported for {0}r'   r(   )r   +statsmodels.genmod.generalized_linear_modelrI   #statsmodels.discrete.discrete_modelrK   getattr
isinstanceskipformattypefittedvaluesr   modelZendogZresidpredict)r7   r   rI   rK   r8   fittedr   r   r   check_fittedx   s    rX   c                 C   s  | }t t |jjdd }ddlm} ddlm} ddl	m
}m} t| d| } t| ||fr|| ||  ||d   n|jdd }t|||dd	 t||t | dd	 t|dd
 ||d  dd	 t|dd
 ||d dd	 tt|}||}	|jd
kr<tjntj}
||
||d}|	jd
krftjntj}
|
|	|d}t|tjr||| n
||| dS )z
    Check that the `predict` method of the given results object produces the
    correct output type.

    Parameters
    ----------
    results : Results

    Raises
    ------
    AssertionError
    N   r   rH   rJ   )assert_frame_equalassert_series_equalrL   r'   r(   r   )index)r,   r0   ZasarrayrU   ZexogrM   rI   rN   rK   Zstatsmodels.compat.pandasrZ   r[   rO   rP   rV   tolistrT   r   rD   r.   ndimpdZSeriesZ	DataFrame)r7   r8   Zp_exogrI   rK   rZ   r[   rW   Z
exog_indexZ	predictedclsZpredicted_pandasZpredicted_expectedr   r   r   check_predict_types   s<    

ra   )__doc__r   r
   Zpackaging.versionr   r   numpyr,   Znumpy.testingr   r   Zpandasr_   r   r;   rG   rX   ra   r   r   r   r   <module>   s   ,