U
    *ifD                     @   s  d dl m Z  d dlZd dlZd dlm  mZ d dlm	Z	m
Z
mZmZmZmZ d dlmZ d dlm  mZ ejdejdi gejdi gejdd d	i ejd
dejdi gejdi gejdi gdd di gdd ddd igdd di gdd ddd igg
dd Zejdejdi gejdi gejdd d	i ejd
dejdi gejdi gejdi gdd di gdd ddd igdd di gdd ddd igg
dd Zejdejdi d gejdi d gd!d d	i d gejdi d gejdi d gejdi d gd"d di d gd#d ddd id gd$d di d gd%d ddd id gg
d&d' Z ejdejdi d gejdi d gd(d d	i d gejdi d gejdi d gejdi d gd)d di d gd*d ddd id gd+d di d gd,d ddd id gg
d-d. Z!ejdejdi gejdi gejdi gejdi gejdi gd/d di gd0d ddd igd1d di gd2d ddd igg	d3d4 Z"d5d6 Z#ejd7di gdi gdi gdi gdi gdi gddd igdi gddd igg	ejd8d d9d:gd;d< Z$d=d> Z%ejd?di d@gdi d@gd	i d gdi d@gdi d@gdi d@gdi d@gddd id@gdi d@gddd id@gg
dAdB Z&ejdCdi d dgdi d dgd	i d d gdi d dgdi d dgdi d dgdi d dgddd id dgdi d dgddd id dgg
dDdE Z'ejdCdi d dgdi d dgd	i d d gdi d dgdi d dgdi d dgdi d dgddd id dgdi d dgddd id dgg
dFdG Z(ejdHdId dJd dKd dLd dMd dNd dOd dPd dQd dRd dSd dTd dUd dVd ejdWd ej)dgdXdY Z*dZd[ Z+d\d] Z,d^d_ Z-d`da Z.dbdc Z/ejdddedf dgD dhdf diD  djdk Z0ejdHdld dmd dnd dod dpd dqd drd dsd dtd dud dvd dwd dxd dyd dzd ejd{d ej)dgd|d} Z1dS )~    datetimeN)	DataFrameDatetimeIndexSeriesconcatisnanotnazcompare_func, roll_func, kwargsmeansumc                 C   s   t | t S Nnpisfiniteastypefloatr   x r   ]/home/mars/bis/venv/lib/python3.8/site-packages/pandas/tests/window/test_rolling_functions.py<lambda>       r   countz ignore:min_periods:FutureWarning)Zmarksmedianminmaxc                 C   s   t j| ddS N   ddofr   stdr   r   r   r   r   #   r   r!   c                 C   s   t j| ddS Nr   r   r    r   r   r   r   r   $   r   r   c                 C   s   t j| ddS r   r   varr   r   r   r   r   %   r   r$   c                 C   s   t j| ddS r"   r#   r   r   r   r   r   &   r   c                 C   sF   t | d|f |}t|ts$tt|jd || dd   d S )N2   )getattrrolling
isinstancer   AssertionErrortmassert_almost_equaliloc)seriescompare_func	roll_funckwargsresultr   r   r   test_series   s    r4   c                 C   s   t | t S r   r   r   r   r   r   r   5   r   c                 C   s   t j| ddS r   r    r   r   r   r   r   =   r   c                 C   s   t j| ddS r"   r    r   r   r   r   r   >   r   c                 C   s   t j| ddS r   r#   r   r   r   r   r   ?   r   c                 C   s   t j| ddS r"   r#   r   r   r   r   r   @   r   c                 C   sd   t |d|f |}t|ts$ttj|jdd d f |jdd d d f j|d| ddd d S )Nr%   r&   r'   r   )ZaxisrawFZcheck_names)	r(   r)   r*   r   r+   r,   assert_series_equalr.   apply)r5   framer0   r1   r2   r3   r   r   r   
test_frame/   s     r:   z%compare_func, roll_func, kwargs, minp
   c                 C   s   t | t S r   r   r   r   r   r   r   R   r   c                 C   s   t j| ddS r   r    r   r   r   r   r   V   r   c                 C   s   t j| ddS r"   r    r   r   r   r   r   W   r   c                 C   s   t j| ddS r   r#   r   r   r   r   r   X   r   c                 C   s   t j| ddS r"   r#   r   r   r   r   r   Y   r   c                 C   s~   d}| d d d  d }t|j||d|f |}|jd }|dt   }	| d d d |	|}
t	|d ||
 d S )N      Bwindowmin_periodsr&      )
resampler
   r(   r)   indexoffsetsBDaytruncater,   r-   )r/   r0   r1   r2   minpwinZserZseries_result	last_date	prev_dateZtrunc_seriesr   r   r   test_time_rule_seriesM   s    
rL   c                 C   s   t | t S r   r   r   r   r   r   r   n   r   c                 C   s   t j| ddS r   r    r   r   r   r   r   r   r   c                 C   s   t j| ddS r"   r    r   r   r   r   r   s   r   c                 C   s   t j| ddS r   r#   r   r   r   r   r   t   r   c                 C   s   t j| ddS r"   r#   r   r   r   r   r   u   r   c                 C   s   d}|d d d  d }t|j||d|f |}|jd }	|	dt   }
|d d d |
|	}tj	|
|	|j|| ddd	 d S )
Nr<   r=   r>   r?   r&   rB   r5   Fr6   )rC   r
   r(   r)   rD   rE   rF   rG   r,   r7   Zxsr8   )r5   r9   r0   r1   r2   rH   rI   frmZframe_resultrJ   rK   Ztrunc_framer   r   r   test_time_rule_framei   s    
rO   c                 C   s   t j| ddS r   r    r   r   r   r   r      r   c                 C   s   t j| ddS r"   r    r   r   r   r   r      r   c                 C   s   t j| ddS r   r#   r   r   r   r   r      r   c                 C   s   t j| ddS r"   r#   r   r   r   r   r      r   c                 C   sf  t tjd}tj|d d< tj|dd < t|jddd|f |}t|j	d | |dd  t|jddd|f |}t
|j	d	 stt
|j	d
 rtt
|j	d rtt
|j	d stt tjd}t|jddd|f |}t
|j	d stt|j	d st|dkrbt|jddd|f |}t|jddd|f |}t|| d S )Nr%   r;      rA   r&            rB   i         r   r   r   )r   r   randomrandnNaNr(   r)   r,   r-   r.   r   r+   r	   )r0   r1   r2   objr3   obj2Zresult0Zresult1r   r   r   	test_nans   s$    
r_   c                  C   sl   t tjd} tj| d d< tj| dd < | jddd }t|j	d t
| dd t  d S )Nr%   r;   rP   rQ   rR   r&   )r   r   rZ   r[   r\   r)   r   r,   r-   r.   r   r   r   r   )r]   r3   r   r   r   test_nans_count   s     r`   zroll_func, kwargsrH   c   d   c                 C   sv   t | jt| d |d|f |}t | jt| |d|f |}t|}t|t| | }t|| ||  d S )Nr   rR   )r(   r)   lenr   r,   r7   r-   )r/   rH   r1   r2   r3   expectednan_maskr   r   r   test_min_periods   s    rf   c                 C   sf   | j t| d dd }| j t| dd }t|}t|t| | }t|| ||  d S )Nr   r   rR   )r)   rc   r   r   r,   r7   r-   )r/   r3   rd   re   r   r   r   test_min_periods_count   s    rg   zroll_func, kwargs, minprT   c                 C   s   t tjd}tj|d d< tj|dd < t|jd|dd| f |}tt|t tjgd gjd|d| f |dd  jdd	}t	
|| d S )
Nr%   r;   rP   rS   T)rA   center	   rR   )Zdrop)r   r   rZ   r[   r\   r(   r)   r   Zreset_indexr,   r7   )r1   r2   rH   r]   r3   rd   r   r   r   test_center   s"      rj   z#roll_func, kwargs, minp, fill_valuec                 C   s   dd t dD }t| t| j| jd|d|f |d| j}t| jd|dd|f |}|d k	rx||}t	|| d S )	Nc                 S   s   g | ]}d |dqS r   dr   .0r   r   r   r   
<listcomp>  s     z.test_center_reindex_series.<locals>.<listcomp>   r<   r?   Tr@   rA   rh   )
ranger(   reindexlistrD   r)   shiftfillnar,   r7   )r/   r1   r2   rH   
fill_valuesZ	series_xpZ	series_rsr   r   r   test_center_reindex_series   s*     
rz   c                 C   s   dd t dD }t| t| j| jd|d|f |d| j}t| jd|dd|f |}|d k	rx||}t	|| d S )	Nc                 S   s   g | ]}d |dqS rk   r   rm   r   r   r   ro   4  s     z-test_center_reindex_frame.<locals>.<listcomp>rp   r<   r?   rq   Trr   )
rs   r(   rt   ru   rD   r)   rv   rw   r,   assert_frame_equal)r9   r1   r2   rH   rx   ry   Zframe_xpZframe_rsr   r   r   test_center_reindex_frame#  s*     
r|   fc                 C   s   | j dddj| ddS Nr;   rW   r?   F)pairwiser)   Zcovr   r   r   r   r   I  r   c                 C   s   | j dddj| ddS r~   r)   Zcorrr   r   r   r   r   J  r   c                 C   s   | j ddd S Nr;   rW   r?   r)   r   r   r   r   r   r   K  r   c                 C   s   | j ddd S r   r)   r   r   r   r   r   r   L  r   c                 C   s   | j ddd S r   r)   r   r   r   r   r   r   M  r   c                 C   s   | j ddd S r   r)   r
   r   r   r   r   r   N  r   c                 C   s   | j ddd S r   r)   r!   r   r   r   r   r   O  r   c                 C   s   | j ddd S r   r)   r$   r   r   r   r   r   P  r   c                 C   s   | j ddd S r   r)   Zskewr   r   r   r   r   Q  r   c                 C   s   | j ddd S r   r)   Zkurtr   r   r   r   r   R  r   c                 C   s   | j dddjddS )Nr;   rW   r?         ?)quantiler)   r   r   r   r   r   r   S  r   c                 C   s   | j ddd S r   r)   r   r   r   r   r   r   T  r   c                 C   s   | j dddjtddS Nr;   rW   r?   FrM   r)   r8   r   r   r   r   r   r   U  r   c                 C   s   | j dddjtddS Nr;   rW   r?   TrM   r   r   r   r   r   r   V  r   c                 C   s   | j dddd S NZboxcarr;   rW   )Zwin_typer@   rA   r   r   r   r   r   r   X  r   c                 C   s   t td}t tj|jd}tddgddgddgdd	ggd
dgd}ttj|j|jd}| |}t|| | |}t	|| d S )NrY   rD   r   rW   rX   r=   ri   r&   r   Ar>   columns)rD   r   )
r   rs   r   nanrD   r   r   r,   r7   r{   )r}   ry   
s_expectedZdfZdf_expecteds_resultZ	df_resultr   r   r   +test_rolling_functions_window_non_shrinkageF  s    (r   c                  C   s   dd t ddD } | tddddd tt dd| d	}|d
d }| }tdddddgtdd t ddD ddd	}|d j	dd }t
|| dS )z%Replicate result expected in GH #6297c                 S   s   g | ]}t d d|qS   r   r   rn   ir   r   r   ro   m  s     z+test_rolling_max_gh6297.<locals>.<listcomp>r      r   rX   r      r   c                 S   s   t | S r   r   r   r   r   r   r   r  r   z)test_rolling_max_gh6297.<locals>.<lambda>      ?       @g      @      @g      @c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   r   r   ro   x  s     Dfreqr@   N)rs   appendr   r   map
sort_indexr   rC   r   r)   r,   r7   indicesr/   rd   r   r   r   r   test_rolling_max_gh6297k  s    r   c                  C   st  dd t ddD } | tdddd | tdddd ttt ddd	d
g | d}|dd }| }tdddddgtdd t ddD ddd}|d	 j
dd	 }t|| tdddddgtdd t ddD ddd}|d j
dd	 }t|| d}tdddd|gtdd t ddD ddd}|d j
dd	 }t|| d S )Nc                 S   s   g | ]}t d d|qS r   r   r   r   r   r   ro     s     z-test_rolling_max_resample.<locals>.<listcomp>r   r   r   rW   r=   r   r;   rS   r   c                 S   s   t | S r   r   r   r   r   r   r     r   z+test_rolling_max_resample.<locals>.<lambda>        r   r         @g      4@c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   ro     s     r   r   r   g      $@c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   ro     s     g&@c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   ro     s     )rs   r   r   r   ru   r   r   r   rC   r   r)   r,   r7   r   r
   )r   r/   rd   r   vr   r   r   test_rolling_max_resample~  s2    r   c                  C   s   dd t ddD } | tdddd | tdddd ttt ddd	d
g | d}|dd }| }tdddddgtdd t ddD ddd}|d	 j
dd}t||	  d S )Nc                 S   s   g | ]}t d d|qS r   r   r   r   r   r   ro     s     z-test_rolling_min_resample.<locals>.<listcomp>r   r   r   rW   r=   r   r;   rS   r   c                 S   s   t | S r   r   r   r   r   r   r     r   z+test_rolling_min_resample.<locals>.<lambda>r   r   r   r   r   c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   ro     s     r   r   r   )rs   r   r   r   ru   r   r   r   rC   r   r)   r,   r7   )r   r/   rd   rr   r   r   test_rolling_min_resample  s    r   c                  C   s   dd t ddD } | tdddd | tdddd ttt ddd	d
g | d}|dd }| }tddddd	gtdd t ddD ddd}|d	 j
dd	 }t|| d S )Nc                 S   s   g | ]}t d d|qS r   r   r   r   r   r   ro     s     z0test_rolling_median_resample.<locals>.<listcomp>r   r   r   rW   r=   r   r;   rS   r   c                 S   s   t | S r   r   r   r   r   r   r     r   z.test_rolling_median_resample.<locals>.<lambda>r   r   r   r   c                 S   s   g | ]}t d d|dqS r   r   r   r   r   r   ro     s     r   r   r   )rs   r   r   r   ru   r   r   r   rC   r   r)   r,   r7   r   r   r   r   test_rolling_median_resample  s    r   c                  C   sD   d} t tj| jddd  t tj| jddd  d S )Ni N  r=   F)r@   rh   )r   r   rZ   r[   r)   r   )nr   r   r    test_rolling_median_memory_error  s    r   	data_typec                 C   s   g | ]}t d | qS )r}   r   dtype)rn   widthr   r   r   ro     s     ro   )rY      c                 C   s(   g | ] }d D ]}t | | qqS )Zuir   )rn   r   signr   r   r   ro     s       )r   r=   rY   r   c                 C   sp   t tjd| djdd }|jd tdks6tt tjd| djdd }|jd tdksltd S )NrS   r   rW   r   r   Zf8)	r   r   Zaranger)   r   Zdtypesr   r+   r   )r   r3   r   r   r   "test_rolling_min_max_numeric_types  s    r   c                 C   s   | j ddd S )Nr;   r   r?   )r)   r   r   r   r   r   r     r   c                 C   s   | j dddj| ddS r~   r   r   r   r   r   r     r   c                 C   s   | j dddj| ddS r~   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j ddddS )Nr;   rW   r?   r   r   r   r   r   r   r     r   c                 C   s   | j ddd S r   r   r   r   r   r   r     r   c                 C   s   | j dddjtddS r   r   r   r   r   r   r     r   c                 C   s   | j dddjtddS r   r   r   r   r   r   r     r   c                 C   s   | j dddd S r   r   r   r   r   r   r     r   c           
      C   s|   t tjd}|}t }|}tdgd}|d d|d< |}| |}t|| | |}t|| | |}	t|	| d S )Nr   ar   float64)r   r   r   r   r   r,   r7   r{   )
r}   ry   r   Zdf1Zdf1_expectedZdf2Zdf2_expectedr   Z
df1_resultZ
df2_resultr   r   r   !test_moment_functions_zero_length  s    r   )2r   numpyr   ZpytestZpandas.util._test_decoratorsutilZ_test_decoratorstdZpandasr   r   r   r   r   r	   Zpandas._testingZ_testingr,   Zpandas.tseries.offsetsZtseriesrE   markZparametrizer
   Znansumparamfilterwarningsr   r   r   r4   r:   rL   rO   r_   r`   rf   rg   rj   rz   r|   Zskip_if_no_scipyr   r   r   r   r   r   r   r   r   r   r   r   <module>   s   














































&
