U
    Kvfg                     @   s  d Z ddl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m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mZmZ ejejeZG d	d
 d
ZG dd deZG dd deZ G dd deZ!G dd dZ"G dd de"Z#G dd de"Z$G dd de"Z%G dd dZ&G dd de&Z'G dd de&Z(G dd  d e&Z)G d!d" d"Z*G d#d$ d$e*Z+G d%d& d&e*Z,G d'd( d(Z-G d)d* d*e-Z.G d+d, d,Z/G d-d. d.e/Z0G d/d0 d0e/Z1G d1d2 d2e/Z2G d3d4 d4eZ3dOd6d7Z4e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gd?d@ Z7e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdAdB Z8e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdCdD Z9e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdEdF Z:dGdH Z;e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdIdJ Z<e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdKdL Z=e
j56d8d9d:d;dge
j56d<d5d=ge
j56d>d5d=gdMdN Z>dS )Pa  
Tests for smoothing and estimation of unobserved states and disturbances

- Predicted states: :math:`E(\alpha_t | Y_{t-1})`
- Filtered states: :math:`E(\alpha_t | Y_t)`
- Smoothed states: :math:`E(\alpha_t | Y_n)`
- Smoothed disturbances :math:`E(\varepsilon_t | Y_n), E(\eta_t | Y_n)`

Tested against R (FKF, KalmanRun / KalmanSmooth), Stata (sspace), and
MATLAB (ssm toolbox)

Author: Chad Fulton
License: Simplified-BSD
    N)assert_allcloseassert_almost_equalassert_equal)datasets)mlemodelsarimaxvarmax)TVSS)FILTER_UNIVARIATE)SMOOTH_CLASSICALSMOOTH_ALTERNATIVESMOOTH_UNIVARIATEc                   @   sr   e Zd ZedddZdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd ZdS )TestStatesAR3Fc                 O   s  t jtdd}t|| _tjdddd| j_t jtdd}dd	d
dddddddddg}tj|d |d| _	t
j| jd f|dddd|| _|rd| jj_tjd }| jj|dd| _td| jjf| j_td| jjf| j_t| jjD ]d}tj| jjjd d d d |f | jjd|f< tj| jjjd d d d |f | jjd|f< q| jj}| jj}	| jjj}
| jj dd| _!| j!j"t||	 t||
 t| jj#d d S ) Nresultsresults_wpi1_ar3_stata.csv
1960-01-01|   QSstartZperiodsfreqzresults_wpi1_ar3_matlab_ssm.csva1a2a3detP	alphahat1	alphahat2	alphahat3detVepsepsvaretaetavarheadernameswpi)      r   T)orderZsimple_differencinghamilton_representationg?ghc?g?mT?g_?none)Zcov_typer(   r   )Zfilter_timingZ measurement_disturbance_variatesZstate_disturbance_variatesZinitial_state_variates)$ospathjoincurrent_pathpdread_csvstata
date_rangeindex
matlab_ssmr   SARIMAXmodelssmtiming_init_filterednpr_smoothr   zerosnobsdet_predicted_state_covdet_smoothed_state_covrangelinalgdetfilter_resultspredicted_state_covsmoother_resultssmoothed_state_covk_endogk_posdefsimulation_smoothersimsimulatek_statesclsalternate_timingargskwargsr/   Zmatlab_namesparamsir@   rJ   rK    rW   S/tmp/pip-unpacked-wheel-2v6byqio/statsmodels/tsa/statespace/tests/test_smoothing.pysetup_class#   sr    


           


zTestStatesAR3.setup_classc                 C   s.   t | jj jd | jjdd  d d d S )Nr   r(   Zdep1   )r   r   rF   Zpredict	forecastsr4   ilocselfrW   rW   rX   test_predict_obsU   s
     zTestStatesAR3.test_predict_obsc                 C   s*   t | jjjd | jjdd  d d d S )Nr   r(   Zsr1rZ   )r   r   rF   Zstandardized_forecasts_errorr4   r\   r]   rW   rW   rX   test_standardized_residuals[   s
     z)TestStatesAR3.test_standardized_residualsc                 C   sn   t | jjjd d d df j| jjdd  dddg d t | jjjd d d df j| jddd	g d d S )
Nr(   Zsp1Zsp2Zsp3rZ   r   r   r   )r   r   rF   predicted_stateTr4   r\   r7   r]   rW   rW   rX   test_predicted_statesa   s      z#TestStatesAR3.test_predicted_statesc                 C   s   t | jjj| jdg d d S Nr   rZ   r   r   rA   rc   r7   r]   rW   rW   rX   test_predicted_states_covk   s
    
 z'TestStatesAR3.test_predicted_states_covc                 C   s.   t | jjjj| jjdd  dddg d d S )Nr(   Zsf1Zsf2Zsf3rZ   )r   r   rF   Zfiltered_staterc   r4   r\   r]   rW   rW   rX   test_filtered_statesq   s
    
 z"TestStatesAR3.test_filtered_statesc                 C   sN   t | jjjj| jjdd  dddg d t | jjjj| jdddg d d S )	Nr(   sm1sm2sm3rZ   r   r   r   r   r   rH   smoothed_staterc   r4   r\   r7   r]   rW   rW   rX   test_smoothed_statesw   s    
 
 z"TestStatesAR3.test_smoothed_statesc                 C   s   t | jjj| jdg d d S Nr   rZ   r   r   rB   rc   r7   r]   rW   rW   rX   test_smoothed_states_cov   s
    
 z&TestStatesAR3.test_smoothed_states_covc                 C   s    t | jjjj| jdg d d S Nr   rZ   )r   r   rH    smoothed_measurement_disturbancerc   r7   r]   rW   rW   rX   %test_smoothed_measurement_disturbance   s
    

 z3TestStatesAR3.test_smoothed_measurement_disturbancec                 C   s(   | j j}t|jd j| jdg d d S Nr   r    rZ   )r   rH   r   $smoothed_measurement_disturbance_covrc   r7   )r^   resrW   rW   rX   )test_smoothed_measurement_disturbance_cov   s    

 z7TestStatesAR3.test_smoothed_measurement_disturbance_covc                 C   s    t | jjjj| jdg d d S )Nr!   rZ   )r   r   rH   smoothed_state_disturbancerc   r7   r]   rW   rW   rX   test_smoothed_state_disturbance   s
    

 z-TestStatesAR3.test_smoothed_state_disturbancec                 C   s$   t | jjjd j| jdg d d S )Nr   r"   rZ   )r   r   rH   smoothed_state_disturbance_covrc   r7   r]   rW   rW   rX   #test_smoothed_state_disturbance_cov   s
    
 z1TestStatesAR3.test_smoothed_state_disturbance_covN)F)__name__
__module____qualname__classmethodrY   r_   r`   rd   rg   rh   rn   rq   rt   rx   rz   r|   rW   rW   rW   rX   r   "   s   1

r   c                       s    e Zd Ze fddZ  ZS )TestStatesAR3AlternateTimingc                    s   t t| j|ddi| d S NrR   T)superr   rY   rQ   rS   rT   	__class__rW   rX   rY      s    
z(TestStatesAR3AlternateTiming.setup_classr}   r~   r   r   rY   __classcell__rW   rW   r   rX   r      s   r   c                       s8   e Zd Ze fddZdd Zdd Zdd Z  ZS )	!TestStatesAR3AlternativeSmoothingc                    s   t t| j|dti| d S Nsmooth_method)r   r   rY   r   r   r   rW   rX   rY      s    
z-TestStatesAR3AlternativeSmoothing.setup_classc                 C   sh   t | jjjjdd  | jjdd  dddg d t | jjjjdd  | jjdd  ddd	g d d S )
N   r'   ri   rj   rk   rZ   r   r   r   rl   r]   rW   rW   rX   rn      s     z6TestStatesAR3AlternativeSmoothing.test_smoothed_statesc                 C   s0   t | jjjdd  | jjdd  dg d d S )Nr(   r   rZ   )r   r   rB   rc   r7   r\   r]   rW   rW   rX   rq      s
     z:TestStatesAR3AlternativeSmoothing.test_smoothed_states_covc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S Nr   r9   r:   r   r   _kalman_smoother_smooth_methodr]   rW   rW   rX   test_smooth_method   s    z4TestStatesAR3AlternativeSmoothing.test_smooth_method)	r}   r~   r   r   rY   rn   rq   r   r   rW   rW   r   rX   r      s
   r   c                       s(   e Zd Ze fddZdd Z  ZS ) TestStatesAR3UnivariateSmoothingc                    s   t t| j|dti| d S Nfilter_method)r   r   rY   r
   r   r   rW   rX   rY      s    
z,TestStatesAR3UnivariateSmoothing.setup_classc                 C   s8   t | jjjd t | jjjjd t | jjjjt d S Nr   r   r9   r:   r   r   r   r   r]   rW   rW   rX   r      s
    z3TestStatesAR3UnivariateSmoothing.test_smooth_methodr}   r~   r   r   rY   r   r   rW   rW   r   rX   r      s   r   c                   @   sZ   e Zd ZedddZdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd ZdS )TestStatesMissingAR3Fc                 O   s  t jtdd}t|| _tjdddd| j_t jtdd}dd	d
dddddddddg}tj|d |d| _	t jtdd}t|| _
| jd  | jd< tj| jj| jjdd df< tj| jj| jjdd  df f|ddd|| _|rd| jj_tjd }| jj|dd| _td| jjf| j_td| jjf| j_t| jjD ]b}tj| jjd d d d |f | jjd |f< tj| jjd d d d |f | jjd |f< qV| jj}| jj }	| jjj!}
| j" | _#| j#j$t||	 t||
 t| jj%d! d S )"Nr   r   r   r   r   r   z'results_wpi1_missing_ar3_matlab_ssm.csvr   r   r   r   r   r   r   r   r   r    r!   r"   r#   zresults_smoothing3_R.csvr&   Zdwpi
      r(   )r'   r   r   T)r)   r*   r+   
return_ssmr   r-   )&r.   r/   r0   r1   r2   r3   r4   r5   r6   r7   R_ssmdiffr<   nanlocr   r8   r9   r:   r;   r=   r>   r   r?   r@   rA   rB   rC   rD   rE   rG   rI   rJ   rK   rL   rM   rN   rO   rP   rW   rW   rX   rY      s|    


          



z TestStatesMissingAR3.setup_classc                 C   s2   t | jjd d d df j| jdddg d d S )Nra   r   r   r   rZ   )r   r   rb   rc   r7   r]   rW   rW   rX   rd     s
     z*TestStatesMissingAR3.test_predicted_statesc                 C   s   t | jjj| jdg d d S re   rf   r]   rW   rW   rX   rg     s
    
 z.TestStatesMissingAR3.test_predicted_states_covc                 C   s"   t | jjj| jdddg d d S )Nr   r   r   rZ   )r   r   rm   rc   r7   r]   rW   rW   rX   rn     s
     z)TestStatesMissingAR3.test_smoothed_statesc                 C   s   t | jjj| jdg d d S ro   rp   r]   rW   rW   rX   rq   !  s
    
 z-TestStatesMissingAR3.test_smoothed_states_covc                 C   s   t | jjj| jdg d d S rr   )r   r   rs   rc   r7   r]   rW   rW   rX   rt   '  s
    
 z:TestStatesMissingAR3.test_smoothed_measurement_disturbancec                 C   s"   t | jjd j| jdg d d S ru   )r   r   rv   rc   r7   r]   rW   rW   rX   rx   -  s
    
 z>TestStatesMissingAR3.test_smoothed_measurement_disturbance_covc                 C   s   t | jjj| jdg d d S )NZetahat	   )r   r   ry   rc   r   r]   rW   rW   rX   rz   7  s
    
 z4TestStatesMissingAR3.test_smoothed_state_disturbancec                 C   s(   t | jjddd d f | jd d d S )Nr   detVetar   )r   r   r{   r   r]   rW   rW   rX   r|   B  s
     z8TestStatesMissingAR3.test_smoothed_state_disturbance_covN)F)r}   r~   r   r   rY   rd   rg   rn   rq   rt   rx   rz   r|   rW   rW   rW   rX   r      s   8
r   c                       s    e Zd Ze fddZ  ZS )#TestStatesMissingAR3AlternateTimingc                    s   t t| j|ddi| d S r   )r   r   rY   r   r   rW   rX   rY   O  s    z/TestStatesMissingAR3AlternateTiming.setup_classr   rW   rW   r   rX   r   N  s   r   c                       s(   e Zd Ze fddZdd Z  ZS )(TestStatesMissingAR3AlternativeSmoothingc                    s   t t| j|dti| d S r   )r   r   rY   r   r   r   rW   rX   rY   V  s    
z4TestStatesMissingAR3AlternativeSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r   [  s    z;TestStatesMissingAR3AlternativeSmoothing.test_smooth_methodr   rW   rW   r   rX   r   U  s   r   c                       s(   e Zd Ze fddZdd Z  ZS )'TestStatesMissingAR3UnivariateSmoothingc                    s   t t| j|dti| d S r   )r   r   rY   r
   r   r   rW   rX   rY   d  s    
z3TestStatesMissingAR3UnivariateSmoothing.setup_classc                 C   s8   t | jjjd t | jjjjd t | jjjjt d S r   r   r]   rW   rW   rX   r   i  s
    z:TestStatesMissingAR3UnivariateSmoothing.test_smooth_methodr   rW   rW   r   rX   r   c  s   r   c                   @   s   e Zd ZdZedd Zdd Zdd Zdd	 Zd
d Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"S )#TestMultivariateMissinga  
    Tests for most filtering and smoothing variables against output from the
    R library KFAS.

    Note that KFAS uses the univariate approach which generally will result in
    different predicted values and covariance matrices associated with the
    measurement equation (e.g. forecasts, etc.). In this case, although the
    model is multivariate, each of the series is truly independent so the
    values will be the same regardless of whether the univariate approach
    is used or not.
    c                 K   sd  t jtdd}t|| _tj	 j
}tjdddd|_|ddd	g  jd
d  }tj|jdddf< tj|jddd
f< tj|jdddf< tj|jdddf< tj|jdddf< tj|fddd|}td|d< td|d< td|d< td|d< td|d< |d || _|jg dd| _td
| jjf| j_td
| jjf| j_td
| jjf| j_td
| jjf| j_t| jjD ]}tj | jj!d d d d |f | jjd|f< tj | jj"d d d d |d
 f | jjd|f< tj | jj#d d d d |f | jjd|f< tj | jj$d d d d |f | jjd|f< qd S )Nr   zresults_smoothing_R.csv
1959-01-01	2009-7-01r   r   endr   realgdprealconsrealinvr(   r   2      F   '   Z   r   w      r'   rO   rK   designobs_cov
transition	selection	state_cov    .ATr   )%r.   r/   r0   r1   r2   r3   desiredr   	macrodataload_pandasdatar5   r6   r   r\   r<   r   r   MLEModeleyeinitialize_approximate_diffuser9   r>   r   r?   r@   !det_scaled_smoothed_estimator_covrA   rB   "det_smoothed_state_disturbance_covrC   rD   rE   scaled_smoothed_estimator_covrG   rI   r{   )rQ   rT   r/   dtaobsmodrV   rW   rW   rX   rY   |  sT    
z#TestMultivariateMissing.setup_classc                 C   s   t t| jjd d S )NgH	r   r<   sumr   Zllf_obsr]   rW   rW   rX   test_loglike  s    z$TestMultivariateMissing.test_loglikec                 C   s    t | jjj| jdddg  d S )Nr1r2r3r   r   Zscaled_smoothed_estimatorrc   r   r]   rW   rW   rX   test_scaled_smoothed_estimator  s    z6TestMultivariateMissing.test_scaled_smoothed_estimatorc                 C   s   t | jjj| jdg  d S )NdetN)r   r   r   rc   r   r]   rW   rW   rX   "test_scaled_smoothed_estimator_cov  s    
z:TestMultivariateMissing.test_scaled_smoothed_estimator_covc                 C   s    t | jjj| jdddg  d S )Nm1m2m3r   r   r[   rc   r   r]   rW   rW   rX   test_forecasts  s    z&TestMultivariateMissing.test_forecastsc                 C   s    t | jjj| jdddg  d S )Nv1v2v3r   r   Zforecasts_errorrc   r   r]   rW   rW   rX   test_forecasts_error  s    z,TestMultivariateMissing.test_forecasts_errorc                 C   s"   t | jj | jdddg  d S NF1ZF2ZF3r   r   Zforecasts_error_covdiagonalr   r]   rW   rW   rX   test_forecasts_error_cov  s    
z0TestMultivariateMissing.test_forecasts_error_covc                 C   s0   t | jjd d dd f j| jdddg  d S )Nr(   r   r   r   r   r   rb   rc   r   r]   rW   rW   rX   rd     s    z-TestMultivariateMissing.test_predicted_statesc                 C   s   t | jjj| jdg  d S )Nr   r   r   rA   rc   r   r]   rW   rW   rX   rg     s    
z1TestMultivariateMissing.test_predicted_states_covc                 C   s    t | jjj| jdddg  d S )Nr   r   r   r   r   rm   rc   r   r]   rW   rW   rX   rn     s    z,TestMultivariateMissing.test_smoothed_statesc                 C   s   t | jjj| jdg  d S )Nr   r   r   rB   rc   r   r]   rW   rW   rX   rq     s    
z0TestMultivariateMissing.test_smoothed_states_covc                 C   s    t | jjj| jdddg  d S )Nmuhat1muhat2muhat3r   r   smoothed_forecastsrc   r   r]   rW   rW   rX   test_smoothed_forecasts  s    z/TestMultivariateMissing.test_smoothed_forecastsc                 C   s    t | jjj| jdddg  d S )Netahat1etahat2etahat3r   r   ry   rc   r   r]   rW   rW   rX   rz     s    z7TestMultivariateMissing.test_smoothed_state_disturbancec                 C   s   t | jjj| jdg  d S )Nr   r   r   r   rc   r   r]   rW   rW   rX   r|     s    
z;TestMultivariateMissing.test_smoothed_state_disturbance_covc                 C   s    t | jjj| jdddg  d S )Nepshat1epshat2epshat3r   r   rs   rc   r   r]   rW   rW   rX   rt     s    z=TestMultivariateMissing.test_smoothed_measurement_disturbancec                 C   s"   t | jj | jdddg  d S NVeps1Veps2Veps3r   r   rv   r   r   r]   rW   rW   rX   rx     s    
zATestMultivariateMissing.test_smoothed_measurement_disturbance_covNr}   r~   r   __doc__r   rY   r   r   r   r   r   r   rd   rg   rn   rq   r   rz   r|   rt   rx   rW   rW   rW   rX   r   p  s$   
0r   c                       s(   e Zd Ze fddZdd Z  ZS ))TestMultivariateMissingClassicalSmoothingc                    s   t t| j|dti| d S r   )r   r   rY   r   r   r   rW   rX   rY     s    
z5TestMultivariateMissingClassicalSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r9   r:   r   r   r   r   r]   rW   rW   rX   r     s    z<TestMultivariateMissingClassicalSmoothing.test_smooth_methodr   rW   rW   r   rX   r     s   r   c                       s(   e Zd Ze fddZdd Z  ZS )+TestMultivariateMissingAlternativeSmoothingc                    s   t t| j|dti| d S r   )r   r   rY   r   r   r   rW   rX   rY     s    
z7TestMultivariateMissingAlternativeSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r     s    z>TestMultivariateMissingAlternativeSmoothing.test_smooth_methodr   rW   rW   r   rX   r     s   r   c                       s(   e Zd Ze fddZdd Z  ZS )*TestMultivariateMissingUnivariateSmoothingc                    s   t t| j|dti| d S r   )r   r   rY   r
   r   r   rW   rX   rY   "  s    
z6TestMultivariateMissingUnivariateSmoothing.setup_classc                 C   s8   t | jjjd t | jjjjd t | jjjjt d S r   r   r]   rW   rW   rX   r   '  s
    z=TestMultivariateMissingUnivariateSmoothing.test_smooth_methodr   rW   rW   r   rX   r   !  s   r   c                   @   s   e Zd ZdZedd Zdd Zdd Zdd	 Zd
d Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"S )#TestMultivariateVAR  
    Tests for most filtering and smoothing variables against output from the
    R library KFAS.

    Note that KFAS uses the univariate approach which generally will result in
    different predicted values and covariance matrices associated with the
    measurement equation (e.g. forecasts, etc.). In this case, although the
    model is multivariate, each of the series is truly independent so the
    values will be the same regardless of whether the univariate approach is
    used or not.
    c                 O   sN  t jtdd}t|| _tj	 j
}tjdddd|_t|ddd	g  jd
d  }tj|fddd|}td|d< tdddgdddgdddgg|d< tdddgdddgdddgg|d< td|d< tdddgddd gdd dgg|d!< |d" || _|jg d#d$| _td
| jjf| j_td
| jjf| j_td
| jjf| j_td
| jjf| j_t| jjD ]}tj !| jj"d d d d |f | jjd%|f< tj !| jj#d d d d |d
 f | jjd%|f< tj !| jj$d d d d |f | jjd%|f< tj !| jj%d d d d |f | jjd%|f< qd S )&Nr   results_smoothing2_R.csvr   r   r   r   r   r   r   r(   r'   r   r   R?        ڇ?Sx[?r   A+o?#c
r?^?'+?K~z\G	$|J@A[?r   r   yN^?P.*,?ZȔ>r   r   Tr   r   )&r.   r/   r0   r1   r2   r3   r   r   r   r   r   r5   r6   r<   logr   r\   r   r   r   arrayr   r9   r>   r   r?   r@   r   rA   rB   r   rC   rD   rE   r   rG   rI   r{   rQ   rS   rT   r/   r   r   r   rV   rW   rW   rX   rY   :  sf    
"



zTestMultivariateVAR.setup_classc                 C   s   t t| jjd d S Nge}@r   r]   rW   rW   rX   r   p  s    z TestMultivariateVAR.test_loglikec                 C   s$   t | jjj| jdddg dd d S Nr   r   r   -C6?atolr   r]   rW   rW   rX   r   s  s
     z2TestMultivariateVAR.test_scaled_smoothed_estimatorc                 C   s<   t t| jjjd d t| jdg d d dd d S )Nra   r   ư>r  r   r<   r  r   r   rc   r   r]   rW   rW   rX   r   y  s
     z6TestMultivariateVAR.test_scaled_smoothed_estimator_covc                 C   s$   t | jjj| jdddg dd d S )Nr   r   r   r  r  r   r]   rW   rW   rX   r     s
     z"TestMultivariateVAR.test_forecastsc                 C   s*   t | jjjd d df | jd dd d S )Nr   r   r  r  r   r]   rW   rW   rX   r     s
     z(TestMultivariateVAR.test_forecasts_errorc                 C   s,   t | jj d d df | jd dd d S )Nr   r   r  r  r   r]   rW   rW   rX   r     s
     z,TestMultivariateVAR.test_forecasts_error_covc                 C   s4   t | jjd d dd f j| jdddg dd d S )Nr(   r   r   r   r  r  r   r]   rW   rW   rX   rd     s
     z)TestMultivariateVAR.test_predicted_statesc                 C   s    t | jjj| jdg dd d S )Nr   缉ؗҜ<r  r   r]   rW   rW   rX   rg     s
    
 z-TestMultivariateVAR.test_predicted_states_covc                 C   s$   t | jjj| jdddg dd d S Nr   r   r   r  r  r   r]   rW   rW   rX   rn     s
     z(TestMultivariateVAR.test_smoothed_statesc                 C   s    t | jjj| jdg dd d S )Nr   r  r  r   r]   rW   rW   rX   rq     s
    
 z,TestMultivariateVAR.test_smoothed_states_covc                 C   s$   t | jjj| jdddg dd d S Nr   r   r   r  r  r   r]   rW   rW   rX   r     s
     z+TestMultivariateVAR.test_smoothed_forecastsc                 C   s$   t | jjj| jdddg dd d S Nr   r   r   r  r  r   r]   rW   rW   rX   rz     s
     z3TestMultivariateVAR.test_smoothed_state_disturbancec                 C   s    t | jjj| jdg dd d S Nr   C]r2<r  r   r]   rW   rW   rX   r|     s
    
 z7TestMultivariateVAR.test_smoothed_state_disturbance_covc                 C   s$   t | jjj| jdddg dd d S Nr   r   r   r  r  r   r]   rW   rW   rX   rt     s
     z9TestMultivariateVAR.test_smoothed_measurement_disturbancec                 C   s&   t | jj | jdddg dd d S )Nr   r   r   r  r  r   r]   rW   rW   rX   rx     s
    
 z=TestMultivariateVAR.test_smoothed_measurement_disturbance_covNr   rW   rW   rW   rX   r   .  s$   
5r   c                       s(   e Zd Ze fddZdd Z  ZS )'TestMultivariateVARAlternativeSmoothingc                    s   t t| j|dti| d S r   )r   r$  rY   r   r   r   rW   rX   rY     s    
z3TestMultivariateVARAlternativeSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r     s    z:TestMultivariateVARAlternativeSmoothing.test_smooth_methodr   rW   rW   r   rX   r$    s   r$  c                       s(   e Zd Ze fddZdd Z  ZS )%TestMultivariateVARClassicalSmoothingc                    s   t t| j|dti| d S r   )r   r%  rY   r   r   r   rW   rX   rY     s    
z1TestMultivariateVARClassicalSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r     s    z8TestMultivariateVARClassicalSmoothing.test_smooth_methodr   rW   rW   r   rX   r%    s   r%  c                   @   s   e Zd ZdZedd Zdd Zdd Zdd	 Zd
d Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"S )#TestMultivariateVARUnivariater  c                 O   sV  t jtdd}t|| _tj	 j
}tjdddd|_t|ddd	g  jd
d  }tj|fddd|}d|j_td|d< tdddgdddgdddgg|d< tdddgdddgdddgg|d< td|d< tddd gddd!gd d!dgg|d"< |d# || _|jg dd$| _td
| jjf| j_td
| jjf| j_td
| jjf| j_td
| jjf| j_ t!| jjD ]}tj"#| jj$d d d d |f | jjd%|f< tj"#| jj%d d d d |d
 f | jjd%|f< tj"#| jj&d d d d |f | jjd%|f< tj"#| jj'd d d d |f | jj d%|f< qd S )&Nr   r  r   r   r   r   r   r   r   r(   r'   r   Tr   r  r  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   r3   r   r   r   r   r   r5   r6   r<   r  r   r\   r   r   r:   filter_univariater   r  r   r9   r>   r   r?   r@   r   rA   rB   r   rC   rD   rE   r   rG   rI   r{   r  rW   rW   rX   rY     sh    
"



z)TestMultivariateVARUnivariate.setup_classc                 C   s   t t| jjd d S r  r   r]   rW   rW   rX   r   (  s    z*TestMultivariateVARUnivariate.test_loglikec                 C   s$   t | jjj| jdddg dd d S r  r   r]   rW   rW   rX   r   +  s
     z<TestMultivariateVARUnivariate.test_scaled_smoothed_estimatorc                 C   s8   t t| jjjd d t| jdg d d  d S )Nra   r   r  r]   rW   rW   rX   r   1  s    z@TestMultivariateVARUnivariate.test_scaled_smoothed_estimator_covc                 C   s*   t | jjjd d df | jd dd d S )Nr   r   r  r  r   r]   rW   rW   rX   r   8  s
     z,TestMultivariateVARUnivariate.test_forecastsc                 C   s$   t | jjj| jdddg dd d S )Nr   r   r   r  r  r   r]   rW   rW   rX   r   >  s
     z2TestMultivariateVARUnivariate.test_forecasts_errorc                 C   s"   t | jj | jdddg  d S r   r   r]   rW   rW   rX   r   D  s    
z6TestMultivariateVARUnivariate.test_forecasts_error_covc                 C   s4   t | jjd d dd f j| jdddg dd d S )Nr(   r   r   r   g:0yE>r  r   r]   rW   rW   rX   rd   J  s
     z3TestMultivariateVARUnivariate.test_predicted_statesc                 C   s    t | jjj| jdg dd d S )Nr   r"  r  r   r]   rW   rW   rX   rg   P  s
    
 z7TestMultivariateVARUnivariate.test_predicted_states_covc                 C   s$   t | jjj| jdddg dd d S r  r   r]   rW   rW   rX   rn   V  s
     z2TestMultivariateVARUnivariate.test_smoothed_statesc                 C   s    t | jjj| jdg dd d S )Nr   r"  r  r   r]   rW   rW   rX   rq   \  s
    
 z6TestMultivariateVARUnivariate.test_smoothed_states_covc                 C   s$   t | jjj| jdddg dd d S r  r   r]   rW   rW   rX   r   b  s
     z5TestMultivariateVARUnivariate.test_smoothed_forecastsc                 C   s$   t | jjj| jdddg dd d S r   r   r]   rW   rW   rX   rz   h  s
     z=TestMultivariateVARUnivariate.test_smoothed_state_disturbancec                 C   s    t | jjj| jdg dd d S r!  r   r]   rW   rW   rX   r|   n  s
    
 zATestMultivariateVARUnivariate.test_smoothed_state_disturbance_covc                 C   s$   t | jjj| jdddg dd d S r#  r   r]   rW   rW   rX   rt   t  s
     zCTestMultivariateVARUnivariate.test_smoothed_measurement_disturbancec                 C   s"   t | jj | jdddg  d S r   r   r]   rW   rW   rX   rx   z  s    
zGTestMultivariateVARUnivariate.test_smoothed_measurement_disturbance_covNr   rW   rW   rW   rX   r&    s$   
6r&  c                       s0   e Zd Ze fddZdd Zdd Z  ZS )&TestMultivariateVARUnivariateSmoothingc                    s   t t| j|dti| d S r   )r   r(  rY   r
   r   r   rW   rX   rY     s    
z2TestMultivariateVARUnivariateSmoothing.setup_classc                 C   s&   t | jjjt t | jjjjt d S r   r   r9   r:   r   r
   r   r]   rW   rW   rX   test_filter_method  s    z9TestMultivariateVARUnivariateSmoothing.test_filter_methodc                 C   s8   t | jjjd t | jjjjd t | jjjjt d S r   r   r]   rW   rW   rX   r     s
    z9TestMultivariateVARUnivariateSmoothing.test_smooth_methodr}   r~   r   r   rY   r*  r   r   rW   rW   r   rX   r(    s   r(  c                   @   s"   e Zd ZedddZdd ZdS )TestVARAutocovariancesmixedc                 O   s  t j j}tjdddd|_t|dddg 	 j
dd  }|d	krztj|j
d d
d d f< tj|j
ddd d f< n|dkrtj|j
dd
df< tj|j
dddf< nn|dkrtj|j
dd
df< tj|j
dddf< tj|j
dddf< tj|j
dddf< tj|j
dddf< tj|fddd|}td|d< tdddgdddgdddgg|d< tdddgd d!d"gd#d$d%gg|d&< td|d'< td(d)d*gd)d+d,gd*d,d-gg|d.< |d/ || _|jg d0d1| _|d2d  tj|fd3dd|}td|dd dd df< tdddgdddgdddgg|d< tdddgd d!d"gd#d$d%gg|d&d dd df< td|d&dd d df< td|d'd dd df< td(d)d*gd)d+d,gd*d,d-gg|d.< |d/ || _|jg d0d1| _d S )4Nr   r   r   r   r   r   r   r(   allr   r   r   partialr   r-  r   r   r   r   r   r'   r   r   g@r  g=z4
?g%'Q+_@r   g^g
G?g
{vX?gίE.gXOE@gƹ?gja?g4CϿg8ur   r   g1=NC@g̿Q%@gT|l@g+=}@gpqQ@g0Z$@r   r   Tr   Zfilter_collapsed   )r   r   r   r   r2   r5   r6   r<   r  r   r\   r   r   r   r   r  r   r9   r>   r   popZaugmented_modelaugmented_results)rQ   whichrS   rT   r   r   r   rW   rW   rX   rY     sz    "







z"TestVARAutocovariances.setup_classc                 C   s   t | jjd d d d ddf | jjd ddd ddf dd t | jjd d d d ddf | jjd ddd dd f d	d d S )
Nr      r'   r(   r0  r  r  ra   gHz>)r   r   Zsmoothed_state_autocovr2  rI   r]   rW   rW   rX   test_smoothed_state_autocov  s    z2TestVARAutocovariances.test_smoothed_state_autocovN)r-  )r}   r~   r   r   rY   r5  rW   rW   rW   rX   r,    s   ?r,  c                       s(   e Zd Ze fddZdd Z  ZS )*TestVARAutocovariancesAlternativeSmoothingc                    s   t t| j|dti| d S r   )r   r6  rY   r   r   r   rW   rX   rY     s    
z6TestVARAutocovariancesAlternativeSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r     s    z=TestVARAutocovariancesAlternativeSmoothing.test_smooth_methodr   rW   rW   r   rX   r6    s   r6  c                       s(   e Zd Ze fddZdd Z  ZS )(TestVARAutocovariancesClassicalSmoothingc                    s   t t| j|dti| d S r   )r   r7  rY   r   r   r   rW   rX   rY     s    
z4TestVARAutocovariancesClassicalSmoothing.setup_classc                 C   s8   t | jjjt t | jjjjt t | jjjjt d S r   r   r]   rW   rW   rX   r     s    z;TestVARAutocovariancesClassicalSmoothing.test_smooth_methodr   rW   rW   r   rX   r7    s   r7  c                       s0   e Zd Ze fddZdd Zdd Z  ZS ))TestVARAutocovariancesUnivariateSmoothingc                    s   t t| j|dti| d S r   )r   r8  rY   r
   r   r   rW   rX   rY     s    
z5TestVARAutocovariancesUnivariateSmoothing.setup_classc                 C   s&   t | jjjt t | jjjjt d S r   r)  r]   rW   rW   rX   r*    s    z<TestVARAutocovariancesUnivariateSmoothing.test_filter_methodc                 C   s8   t | jjjd t | jjjjd t | jjjjt d S r   r   r]   rW   rW   rX   r     s
    z<TestVARAutocovariancesUnivariateSmoothing.test_smooth_methodr+  rW   rW   r   rX   r8    s   r8  c                       s   e Zd Z fddZ  ZS )TVSSWithLagsc                    s@   t  j|dd tdd | ddd d df< | jd d S )N   )Z	_k_statesr0  ).Nr   r   r  )r   __init__r<   r   r:   r   )r^   endogr   rW   rX   r;    s     zTVSSWithLags.__init__)r}   r~   r   r;  r   rW   rW   r   rX   r9    s   r9  Tc                 C   s  t j j}tjdddd|_t|ddg 	 j
dd  }| dkrxtj|j
d d	d d f< tj|j
d
dd d f< nv| dkrtj|j
dd	df< tj|j
d
ddf< nD| dkrtj|j
dd	df< tj|j
dddf< tj|j
d
ddf< |d k	r tj|jd t|| dd}||}|s`tj|dddd}	||	j_|d krP|	j}|	j||d}
nt|}	||	j_|	jg |d}
|	|
fS )Nr   r   r   r   r   r   r(   r.  r4        r/  r   r-     r   )rZ   r   r   T)r)   Zmeasurement_errorZ	tolerancer   )r   r   r   r   r2   r5   r6   r<   r  r   r\   r   lenZreindexr   ZVARMAXr:   r'  start_paramsr>   r9  )missingr'  tvpoosrU   r   r   r<  Znew_ixr   rw   rW   rW   rX   get_acov_model  sB     

 


rE  rB  r.  r/  r-  r'  FrC  c                 C   s  t | ||\}}|jddd}|ddddddf }|ddddddf }|ddddddf }|dddd}	t|	ddddddf |dd ddd	 t|	dd tj |dddd}
t|
ddddddf |dd ddd	 t|
dd tj |d
ddd}t|d
dddddf |d
d ddd	 t|dd
 tj |jddd}	t|	tj |jddd}	t|	ddddf |d ddd	 |jddddddd}	t|	ddddddf |dd ddd	 |jddd}
t|
tj |jddd}
t|
tj |jddd}
t|
ddddf |d ddd	 |jddddddd}
t|
ddddddf |dd ddd	 dS )z"
    Test for Cov(t, t - lag)
    r   r   r(   NrZ   r0  r:  r  )Zrtolr  r'   tr   r   r   )rE  rI   	transposesmoothed_state_autocovariancer   r   r<   r   )rB  r'  rC  _rw   covdesired_acov1desired_acov2desired_acov3acov1acov2acov3rW   rW   rX   -test_smoothed_state_autocovariances_backwards=  sX    ...$    .$    rS  c                 C   s`  t | ||dd\}}ddddddg}|s@t | |||jd	\}}n:t | ||\}}|D ]}	||	d
ddf ||	< qT|j }i }
i }|rddddddg}|D ]0}||d
ddf |
|< ||d
ddf ||< qt|j|j |jddd}|ddddddf }|ddddddf }|ddddddf }tj	|j|j
d
ddf fddddd}|dddd}t|ddddddf |dd  t|ddddddf |ddddddf  |jd|
dddd}t|ddddddf |dd  t|ddddddf |ddddddf  |jd|dddd}t|ddddddf |dd  t|ddddddf |ddddddf  |jd|j|
d}t|ddddf |dddddf  |jddd}t|ddddf |d  |jddddddd}t|ddddddf |dd  |jd|j|d}t|ddddf |dddddf  |jd|jd |
d}t|ddddf |dddddf  |jddd}t|ddddf |d  |jddddddd}t|ddddddf |dd  dS )z"
    Test for Cov(t, t + lag)
    r'   rD  obs_interceptr   r   r   r   r   rU   .Nra   r   r   r(   rZ   r0  r:  Zaxis)extend_kwargs)rG  r[  rF  r   rH  r   r=  )rE  rA  r:   r>   r   llfrI   rI  r<   concatenaterG   rJ  r@   )rB  r'  rC  mod_oosres_oosr%   r   rw   rK  nameZextend_kwargs1Zextend_kwargs2keyskeyrL  rM  rN  rO  Zoos_covrP  rQ  rR  rW   rW   rX   ,test_smoothed_state_autocovariances_forwardsm  s    

   (6   (6   (6  ,    (  ,  ,    rd  c              
   C   s  t | ||dd\}}ddddddg}|s@t | |||jd	\}}n:t | ||\}}|D ]}	||	d
d df ||	< qT|j }t|j|j tj|j|j	d
dd f fdd
ddd}
|
d d ddd df }|
d d ddd df }|
d d ddd df }i }|rr|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}|jd|j|d
ddd}t|j|j|j|jf t|d d d dd df |dd   |jd|jd |d
ddd}t|j|jd |j|jf t|d d d dd df |dd   |jd|jd |d
ddd}t|j|jd |j|jf t|d d d dd df |dd   d S )Nr4  rT  rU  r   r   r   r   r   rV  .rZ  ra   r   rY  r   r(   rZ   r0  r:  rU  r   r   r   r   r   r   r[  rX  rW  r'   )rE  rA  r:   r>   r   r]  r<   r^  rI   rG   rI  rJ  r@   r   shaperO   )rB  r'  rC  r_  r`  r%   r   rw   rK  ra  rL  rM  rN  rO  r[  rP  rQ  rR  rW   rW   rX   0test_smoothed_state_autocovariances_forwards_oos  s    


      (
   (   rh  c              
   C   s  t | ||dd\}}ddddddg}|s@t | |||jd	\}}n:t | ||\}}|D ]}	||	d
d df ||	< qT|j }t|j|j tj|j|j	d
dd f fdd
ddd}
|
d d d dddf }|
d d d dddf }|
d d d dddf }|jd }i }|r||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}|jd||d
ddd}t|j|jd |j|jf t|dd d dd df |dd   t|d d tj |jd||d
ddd}t|dd d dd df |dd   t|d d tj |jd||d
ddd}t|dd d dd df |dd   t|d d tj d S )Nr4  rT  rU  r   r   r   r   r   rV  .rZ  ra   r   rY  r   r(   rZ   r0  r:  re  rf  r'   )rE  rA  r:   r>   r   r]  r<   r^  rI   rG   rI  r@   rJ  r   rg  rO   r   )rB  r'  rC  r_  r`  r%   r   rw   rK  ra  rL  rM  rN  rO  r   r[  rP  rQ  rR  rW   rW   rX   1test_smoothed_state_autocovariances_backwards_oos  s    


  
    (    (    (ri  c               	   C   s   t dddd\} }tjtdd |jdddd W 5 Q R X tjtdd |jddd	 W 5 Q R X tjtdd |jddd
 W 5 Q R X tjtdd |jddd W 5 Q R X tjtdd |jdddd W 5 Q R X d S )NF)rB  r'  rC  zCannot specify both `t`matchr(   )rG  r   zNegative `t`ra   rF  )r   )r   z`end` must be after `start`r4  rZ   rH  )rE  pytestraises
ValueErrorrJ  )rK  rw   rW   rW   rX   +test_smoothed_state_autocovariances_invalid^  s    ro  c                 C   s2  t | ||\}}|rg n|j}tdtj }ddg|d< tj|j|fdd}||}|j|dd}	|	 }
tj|
dd < ||
}|j|dd}dd	d
|j
d	 |j
|j
d	 |j
d fD ]r}|	j||d}|	jd|f |jd|f  }t|j|dd t|jd  |	j|||d	 d}t|j|d dd qd S )N)r   r   g?gɿr   rY  Tr   r(      r   rF  .g+=r  rH  )N.)rE  rA  r<   r?   r   r^  r<  cloner>   copyr@   newsr   r   update_impactsr   revision_impacts)rB  r'  rC  r   rw   rU   appendendog2mod2res2Zendog3Zmod3Zres3rG  outr   rW   rW   rX   test_news_basicr  s(    

*r|  c                 C   s   t | ||dd\}}|rg n|j}|j }ddg|d< ddg|d< ||}|j|dd	}|jd }	d
dd|	d |	|	d |	d fD ]n}
|j||
d}|jd|
f |j	jd|
f  }t
|j|dd |j	jd|
f |jd|
f  }t
|j|dd qd S )Nr   rT  r  ig333333ӿgٿrp  Tr   r   r(   rq  r   rF  .g|=r  )rE  rA  r<  rs  rr  r>   r@   rt  r   Zrevision_resultsr   ru  rv  )rB  r'  rC  r   rw   rU   rx  ry  rz  r@   rG  r{  r   rW   rW   rX   test_news_revisions  s$    


"r}  c              	   C   s  d}t | ||dd\}}|r g n|j}|j }ddg|d< ||}|j|dd}	|j|dd}
|
j||jd d	 d
}t	j
t|d |	j||jd d	 W 5 Q R X |rd}t	j
t|d |
j||jd d	 W 5 Q R X t | ||\}}|rg n|j}|j d |jd  }||}|j|dd}|d }t	j
t|d |j||jd d	 W 5 Q R X tt|j}|jddgdd}|d }t	j
t|d |j||jd d	 W 5 Q R X t | || dd\}}|rd}nd}t	j
t|d |j||jd d	 W 5 Q R X d S )NzThis results object has %s and so it does not appear to by an extension of `previous`. Can only compute the news by comparing this results set to previous results objects.r(   rT  g?g      ?ra   Tr   rF  zECannot compute news without having applied the Kalman smoother first.rj  z[Cannot compute the impacts of news on periods outside of the sample in time-varying models.r   z"fewer observations than `previous`g      ?z0different state space dimensions than `previous`z/time-invariant design while `previous` does notz-time-varying design while `previous` does not)rE  rA  r<  rs  rr  filterr>   rt  r@   rl  rm  rn  RuntimeErrorr   r8   r<   r?   )rB  r'  rC  Zerror_ssr   rw   rU   rx  ry  Zres2_filteredZres2_smoothedmsgrz  rW   rW   rX   test_news_invalid  sF    


r  )NNT)?r   r.   Znumpyr<   Znumpy.testingr   r   r   Zpandasr2   rl  Zstatsmodelsr   Zstatsmodels.tsa.statespacer   r   r   Z7statsmodels.tsa.statespace.tests.test_impulse_responsesr	   Z(statsmodels.tsa.statespace.kalman_filterr
   Z*statsmodels.tsa.statespace.kalman_smootherr   r   r   r/   dirnameabspath__file__r1   r   r   r   r   r   r   r   r   r   r   r   r   r   r$  r%  r&  r(  r,  r6  r7  r8  r9  rE  markZparametrizerS  rd  rh  ri  ro  r|  r}  r  rW   rW   rW   rX   <module>   s    y   M	  
&-TJJ*"