
    !g^                        d Z ddlZddlZddlZddlmZmZm	Z	m
Z
 ddlmZ ddlmZmZ ddlmZ ej$                  j'                         j(                  Z ej,                  ddd	
      e_        d Zd Zej4                  j7                  dddg      ej4                  j7                  dg d      d               Zej4                  j7                  dddg      ej4                  j7                  dg d      d               Zej4                  j7                  dddg      ej4                  j7                  dg d      d               Zd Zej4                  j7                  dddg      ej4                  j7                  dg d      d               Z d Z!ej4                  j7                  dddg      ej4                  j7                  dg d      d               Z"y)zY
Tests for prediction of state space models

Author: Chad Fulton
License: Simplified-BSD
    N)assert_equalassert_raisesassert_allcloseassert_)datasets)sarimaxvarmax)TVSS1959Q12009Q3Q)startendfreqc                     t        j                  ddd      } t        j                  j	                  d       t        j
                  t        j                  j                  d      | d d 	      }t        j                  |d
      }|j                  |j                        }|j                         }t        t        |      |j                         t        |j                  j                   | d d j                          |j#                         }t        |j                  d   | d          t        j                  |dd      }|j                  |j                        }|j                         }t        |j                  |j$                  d   dz
         t        t        |      |j                         t        |j                  j                   | dd j                          |j#                         }t        |j                  d   | d          t        j                  |ddd      }|j                  |j                        }|j                         }t        |j                  |j$                  d   dz
         t        t        |      |j                         t        |j                  j                   | dd j                          |j#                         }t        |j                  d   | d          y )Nz
1950-01-01   D)r   periodsr   i 
   )size)index   r   r   orderr   r   r   r   T)r   simple_differencingr   )r      r   )r   r   r      )r   seasonal_orderr      )pd
date_rangenprandomseedSeriesnormalr   SARIMAXfilterstart_paramspredictr   lennobsr   valuesforecastshape)r   endogmodrespredfcasts         m/var/www/dash_apps/app1/venv/lib/python3.12/site-packages/statsmodels/tsa/statespace/tests/test_prediction.pytest_predict_datesr9      s   MMbsCEIINN6IIbii&&B&/uSbzBE //%y
1C
**S%%
&C ;;=DTCHH%""E#2J$5$56LLNEQr+ //%yd
KC
**S%%
&C;;=D5;;q>A-.TCHH%""E!BK$6$67LLNEQr+ //%y.24C
**S%%
&C;;=D5;;q>U34TCHH%""E%O$:$:;LLNEQr+    c                     g d} t        j                  | d      }|j                  ddg      }d|j                  _        |j                  ddg      }t        |j                  j                  d       t        |j                  d u        t        |j                  j                  d       t        |j                  d u        t        t        |j                  d	       t        t        |j                  d	       t        |j                  d
      |j                  d
             |j!                  d
      }|j!                  d
      }t        |j#                         |j#                                y )N)      ?333333?g?g333333?r   r   r<         ?T)r      )r   r   r?   )dynamicr   )r   r*   r+   ssmmemory_no_predictedr   predicted_stater2   r   predicted_state_covr   
ValueErrorr-   get_predictionr   r1   get_forecastsummary_frame)r3   r4   res1res2fcast1fcast2s         r8   test_memory_no_predictedrM   E   s    E
//%y
1C::sBi D"&CGG::sBi D %%++V4D  D()))//;D$$,- *dllD9*d114@ DMM"%t}}R'89 r"Fr"FF((*F,@,@,BCr:   use_exogTFtrend)nctc                    t        j                  d      j                  dd      dz  }| rt        j                  d      nd }| rd|dd dg}dg}d	g}dg}g }|d
v r||z  }||z  }| r||z  }||z  }|j	                         }	t         j
                  |	dd  t        j                  |	d||      }
|
j                  |      }|j                         }|j                  }| r|d d nd }t        j                  |d d d||      }|j                  |      }| r|dd  nd }|j                  dd|      }|j                  }|j                  |j                  z   |j                  z   }|D ]"  }t        t        ||      t        ||             $ y )Nd   r   r>          @r      皙?r<   r=   rQ   rR   r   r   rO   exog2   r   c   r   r   r[   )r%   arangereshapeonescopynanr   r*   smoothrF   prediction_resultsrepresentation_attributesfilter_attributessmoother_attributesr   getattr)rN   rO   r3   r[   trend_params	ar_paramsexog_params
var_paramsparamsy1mod1rI   p1pr1x2mod2rJ   x2fp2pr2attrskeys                         r8   !test_concatenated_predict_sarimaxrz   b   s    IIcN""3*S0E#2773<DR5LI%KJF
,
iF+
jF	BvvBstH??2Ye$GD;;vD				B


CcrDB??5":Ye"MD;;vD$rs)TC			1"3		7B


C 	%%


	 

!
!	" 
  >S)73+<=>r:   c                    t        j                  d      j                  dd      dz  }| rt        j                  d      nd }ddg}g d}ddg}g d	}g }|d
v r||z  }||z  }| r||z  }||z  }|j	                         }	t         j
                  |	dd  t        j                  |	d||      }
|
j                  |      }|j                         }|j                  }| r|d d nd }t        j                  |d d d||      }|j                  |      }| r|dd  nd }|j                  dd|      }|j                  }|j                  |j                  z   |j                  z   }|D ]"  }t        t        ||      t        ||             $ y )N   rT   r   r>   rW   皙?r<   g        r}   rU   r>   r   r>   rX   rY   r   r   rZ   r\   r   r]   r^   )r%   r_   r`   ra   rb   rc   r	   VARMAXrd   rF   re   rf   rg   rh   r   ri   )rN   rO   r3   r[   rj   rm   rl   
cov_paramsrn   ro   rp   rI   rq   rr   rs   rt   rJ   ru   rv   rw   rx   ry   s                         r8    test_concatenated_predict_varmaxr      s    IIcN""3*S0E#2773<D:L&Jr(KJF
,
jF+
jF	BvvBstH==6TBD;;vD				B


CcrDB==s6RHD;;vD$rs)TC			1"3		7B


C 	%%


	 

!
!	" 
  >S)73+<=>r:   c                    t        j                  d      j                  dd      t         j                  z  }| rt        j                  d      nd }ddg}g d}ddg}g d	}g }|d
v r||z  }||z  }| r||z  }||z  }| r|d d nd }	t        j                  |d d d||	      }
|
j                  |      }| r|dd nd }|j                  ddd|      }|j                  ddd|      }|j                  ddd|      }t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         t        |j                  d d |j                         t        |j                  d d j                  |j                         |j                  dddd|      }|j                  dddd|      }|j                  dddd|      }t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         | du rV|dk(  rQt        |j                  d d |j                         t        |j                  d d j                  |j                         y t        |j                  d d |
d   z   |j                         t        |j                  d d |
d   z   j                  |j                         y )Nr|   rT   r   rW   r}   r~   r>   rU   r   rX   r\   r   rZ   =   r   <   	predictedr   r   information_setr[   filteredsmoothedTr   r   r   signal_onlyr[   FrP   obs_interceptobs_cov)r%   zerosr`   rc   ra   r	   r   rd   rF   r   predicted_meanvar_pred_meanfittedvaluesTforecasts_error_cov)rN   rO   r3   r[   rj   rm   rl   r   rn   x_fitr4   r5   x_fcastp_predf_preds_predp_signalf_signals_signals                      r8   *test_predicted_filtered_smoothed_with_nansr      s5    HHSM!!#q)BFF2E#2773<D:L&Jr(KJF
,
jF+
jF!D"ItE
--cr
&E
JC
**V
C%d2bk4GR   F R   F R   F
 F))6+@+@AF((&*>*>?F))6+@+@AF((&*>*>?F))#2.0@0@AF(("-//1H1HI!!R$ " H !!R " H !!R " H
 H++X-D-DEH**H,B,BCH++X-D-DEH**H,B,BC5Uc\//4c6F6FG..s355s7N7NO//4s?7KK((	*//4s9~EHH//	1r:   c                    t        t        j                  d      t        j                  z        }|j                  j                  ddgt        j                  d             |j                  g       }t        t        j                  d      t        j                  z        }dD ci c]  }|||   
 }} |j                  dddd	d
|} |j                  ddddd
|} |j                  ddddd
|}t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         t        |j                  |j                         t        |j                  d d |j                         t        |j                  d d j                  |j                          |j                  dddd	dd|}	 |j                  dddddd|}
 |j                  dddddd|}t        |j                  |	j                         t        |j                  |	j                         t        |
j                  |	j                         t        |
j                  |	j                         t        |	j                  d d |d   j                  z   |j                         t        |	j                  d d |d   j                  z   j                  |j                         y c c}w )Nr\   r   r=   皙?r   r   r   r   designr   
transition	selection	state_covr   r   r   r   r   r   r   r   r\   Tr   r   r   r   r   r    )r
   r%   r   rc   rA   initialize_knowneyerd   rF   r   r   r   r   r   r   )reset_randomstater4   r5   mod_oosry   kwargsr   r   r   r   r   r   s               r8   /test_predicted_filtered_smoothed_with_nans_TVSSr     s   
rxx 266)
*CGGc3Z3
**R.C288G$rvv-.G,0 1Cc73< 1F 1  S R
F  S R
F  S R
F
 F))6+@+@AF((&*>*>?F))6+@+@AF((&*>*>?F))#2.0@0@AF(("-//1H1HI!s!! R$
H "s!! R
H "s!! R
H
 H++X-D-DEH**H,B,BCH++X-D-DEH**H,B,BCH++CR033G3I3II$$&X++CR03y>3C3CCFF++-W1s   %K<c                    t        j                  t        ddg         }|dv r |j                         j                  dd  dz  }|dk(  r||j                         z  }| rt        j                  d      nd }| rd|dd	 d
dg}g d}ddg}g d}g }|dv r||z  }||z  }| r||z  }||z  }| r|d d nd }	t        j                  |d d d||	      }
ddg|
d<   t        j                  ddgddgg      |
d<   |
j                  |      }| r|dd nd }|j                  ddd|      }|j                  ddd|      }|j                  ddd |      }|j                  d!|"      }|
d   d d d f   }|
d#   }|
d   d d d d d f   }||j                  z  }||j                  z  }||j                  d$d d%f   z  }t!        |j"                  d d ||z   j$                         t!        |j"                  dd  |j"                         t!        |j"                  d d ||z   j$                         t!        |j"                  dd  |j"                         t!        |j"                  d d ||z   j$                         t!        |j"                  dd  |j"                         |d d d d d f   |j&                  j$                  z  |j$                  d d d d d f   z  }|d d d d d f   |j(                  j$                  z  |j$                  d d d d d f   z  }|d d d d d f   |j*                  d$d d%f   j$                  z  |j$                  d d d d d f   z  }t!        |j,                  d d |j$                  |z   j$                         t!        |j,                  dd  |j,                         t!        |j,                  d d |j$                  |z   j$                         t!        |j,                  dd  |j,                         t!        |j,                  d d |j$                  |z   j$                         t!        |j,                  dd  |j,                         |j                  dddd&|'      }|j                  dddd&|'      }|j                  ddd d&|'      }|j"                  |j$                  z
  }|j,                  j$                  |z
  j$                  }t!        |j"                  d d |j$                         t!        |j"                  dd  |d(d)*       t!        |j"                  d d |j$                         t!        |j"                  dd  |       t!        |j"                  d d |j$                         t!        |j"                  dd  |       t!        |j,                  d d |       t!        |j,                  dd  |       t!        |j,                  d d |       t!        |j,                  dd  |       t!        |j,                  d d |       t!        |j,                  dd  |       y )+NrealgdpcpirP   rQ   r   rT   rP   rU      (   rW   r}   r~   r>   r   rX   r\   r   rZ   r?   g       r   r=   g333333?g333333@r   r   r   r   r   r   r   r   r   )r[   r   .r   Tr   ư>:0yE>rtolatol)r%   logdtadiffilocmeanra   r	   r   arrayrd   rF   rG   smoothed_statefiltered_staterC   r   r   r   smoothed_state_covfiltered_state_covrD   r   )rN   rO   r3   r[   rj   rm   rl   r   rn   r   r4   r5   r   r   r   r   r7   dZHdesired_s_signaldesired_f_signaldesired_p_signaldesired_s_signal_covdesired_f_signal_covdesired_p_signal_covr   r   r   fcast_signalfcast_signal_covs                                  r8   'test_predicted_filtered_smoothed_varmaxr   6  sH    FF3	5)*+E


!!!"%+|#2773<DR:L&Jr(KJF
,
jF+
jF!D"ItE
--cr
&E
JC s8CXXSz #Sz+ ,C	N **V
C%d2bk4GR   F R   F R   F
 Rg.EOQW%AHAIq!Tz"A3---3---3..sCRCx88F))#2.5E1E0H0HIF))"#.0D0DEF))#2.5E1E0H0HIF))"#.0D0DEF))#2.5E1E0H0HIF))"#.0D0DE 	
$1*..000133tQz?B  	
$1*..000133tQz?B  	
$1*//SbS9;;;acc$1*oM F(("-0D0F0F0J/M/MNF((-u/B/BCF(("-0D0F0F0J/M/MNF((-u/B/BCF(("-0D0F0F0J/M/MNF((-u/B/BC!!R$ " H !!R " H !!R " H
 ''!##-L++--144H++CR02B2D2DE$lD H++CR02B2D2DEH++BC0,?H++CR02B2D2DEH++BC0,?H**3B/1EFH**23/1ABH**3B/1EFH**23/1ABH**3B/1EFH**23/1ABr:   c                    t        t        j                  d            }|j                  j	                  ddgt        j
                  d             |j                  g       }t        t        j                  d      t        j                  z        }dD ci c]  }|||   
 }} |j                  dddd	d
|} |j                  ddddd
|} |j                  ddddd
|} |j                  dddd	dd|}	 |j                  dddddd|}
 |j                  dddddd|}|d   j                  dd      d d d d d f   }|d   j                  ddd      }|d   j                  ddd      } |j                  di |}|j                  |d   j                  z
  }|j                  |d   j                  z
  }||j                  j                  d d d d d f   z  }||j                  j                  d d d d d f   z  }||j                   j                  d dd d d f   z  }t#        |j                  d d ||z   d          t#        |j                  dd  |j                         t#        |j                  d d ||z   d          t#        |j                  dd  |j                         t#        |j                  d d ||z   d          t#        |j                  dd  |j                         t#        |j                  d d |d          t#        |j                  dd  |dd       t#        |
j                  d d |d          t#        |
j                  dd  |       t#        |	j                  d d |d          t#        |	j                  dd  |       t%        |j&                        D ]  }t#        |j                  |   ||   |j(                  d|f   z  ||   j                  z  ||   z          t#        |j                  |   ||   |j*                  d|f   z  ||   j                  z  ||   z          t#        |j                  |   ||   |j,                  d|f   z  ||   j                  z  ||   z          t#        |j                  |   ||   |j(                  d|f   z  ||   j                  z         t#        |
j                  |   ||   |j*                  d|f   z  ||   j                  z         t#        |	j                  |   ||   |j,                  d|f   z  ||   j                  z          t#        |j                  dd  |j                         t#        |j                  dd  |j                         t#        |j                  dd  |j                         t#        |j                  dd  |       t#        |
j                  dd  |       t#        |	j                  dd  |       y c c}w )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   rA   r   r   rd   rc   rF   	transposerG   r   r   r   r   r   rC   r   ranger/   r   r   rD   )r   r4   r5   r   ry   r   r   r   r   r   r   r   r   r   r   r7   r   r   r   r   r   rR   s                         r8   %test_predicted_filtered_smoothed_TVSSr     s5   
rxx 
!CGGc3Z3
**R.C288G$rvv-.G,0 1Cc73< 1F 1  S R
F  S R
F  S R
F "s!! R$
H "s!! R
H "s!! R
H
 	O&&q!,Q4Z8AH1a(AI  Aq)AC*6*E'''/*B*D*DDL**WY-?-A-AA3--//1d
;;3--//1d
;;3..00"a>>F))#2.5E1Ev0NOF))"#.0D0DEF))#2.5E1Ev0NOF))"#.0D0DEF))#2.5E1Ev0NOF))"#.0D0DEH++CR02B62JK$lD H++CR02B62JKH++BC0,?H++CR02B62JKH++BC0,?388_ I,,Q/!s55c1f==!F1M	O,,Q/!s55c1f==!F1M	O,,Q/!s66sAv>>1G!A$N	P 	..q1!s55c1f==!F	H..q1!s55c1f==!F	H..q1!s66sAv>>1G	II F((-u/B/BCF((-u/B/BCF((-u/B/BCH**23/1ABH**23/1ABH**23/1ABY1s   W'c                 T   t        j                  t        ddg         }|dv r |j                         j                  dd  dz  }|dk(  r||j                         z  }| rt        j                  d      nd }| rd|dd	 d
dg}g d}ddg}g d}g }|dv r||z  }||z  }| r||z  }||z  }| r|d d nd }	| r|dd nd }
t        j                  |d d d||	      }|j                  |      }| r|d d nd }| r|dd nd }t        j                  |d d d||      }|j                  |      }|j                  ddd|
      }|j                  dd|      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
      }|j                  dd|      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
      }|j                  dd|      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
d      }|j                  dd|d      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
dd       }|j                  dd|dd!      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
dd       }|j                  dd|dd!      }t        |j                  |j                         t        |j                  |j                         |j                  ddd|
dd       }|j                  dd|dd!      }t        |j                  |j                         t        |j                  |j                         y )"Nr   r   r   r   rT   rP   rU   r   r   rW   r}   r~   r>   r   rX   r\   r   r   rZ   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   r   r[   r   r   )r%   r   r   r   r   r   ra   r	   r   r+   rF   r   r   r   )rN   rO   r3   r[   rj   rm   rl   r   rn   x_fit1x_fcast1rp   rI   x_fit2x_fcast2rt   rJ   rq   rv   s                      r8   &test_predicted_filtered_dynamic_varmaxr     s1    FF3	5)*+E


!!!"%+|#2773<DR:L&Jr(KJF
,
jF+
jF #T#2YF&tBr{DH==s6VLD;;vD #T#2YF&tBr{DH==s6VLD;;vD 
		1bbx		HB			1"8		<BB%%r'8'89B$$b&6&67			1bbx		HB			1"8		<BB%%r'8'89B$$b&6&67			2t(		KB			22H		=BB%%r'8'89B$$b&6&67 
		1bbx-7 
 
9B			1"8-7 
 
9BB%%r'8'89B$$b&6&67			1bbx-7 
 
9B			1"8-7 
 
9BB%%r'8'89B$$b&6&67			2t(-7 
 
9B			22H-7 
 
9BB%%r'8'89B$$b&6&67 
		1bbx)- 
 
/B			1"8		NBB%%r'8'89B$$b&6&67			1bbx)- 
 
/B			1"8		NBB%%r'8'89B$$b&6&67			2t()- 
 
/B			22H$		OBB%%r'8'89B$$b&6&67 
		1bbx)-z 
 
KB			1"8-7 
 
9BB%%r'8'89B$$b&6&67			1bbx)-z 
 
KB			1"8-7 
 
9BB%%r'8'89B$$b&6&67			2t()-z 
 
KB			22H$-7 
 
9BB%%r'8'89B$$b&6&67r:   )#__doc__pytestnumpyr%   pandasr#   numpy.testingr   r   r   r   statsmodelsr   statsmodels.tsa.statespacer   r	   7statsmodels.tsa.statespace.tests.test_impulse_responsesr
   	macrodataload_pandasdatar   period_ranger   r9   rM   markparametrizerz   r   r   r   r   r   r   r   r:   r8   <module>r      s      O O   6 H $$&++BOO(sC	*,ZD: dE]3/2&> 3 4&>R dE]3/2$> 3 4$>N dE]3/2F1 3 4F1R2-j dE]3/2hC 3 4hCVRCj eT]3/2q8 3 4q8r:   