U
    Cvf'                    @  sL  d dl mZ d dlZd dlZd dlZd dlZd dlm	Z	 d dl
Zd dl
mZmZmZ d dlmZ d dlmZmZmZmZmZmZmZmZ ejdd Zejd	d
 ZddddZddddZddddZddddZ ddddZ!ddddZ"ddddZ#ddddZ$ej%&deddd d!d"d#gd$ddddddgfged%d$ddd d!d"d#gfid$ddddddgigddd&d'Z'ej%&deddd gd$dddgfged%d$ddd gfid$dddgigddd(d)Z(ddd*d+Z)ddd,d-Z*ddd.d/Z+ddd0d1Z,ddd2d3Z-ej%&d4d5d6gddd7d8Z.ddd9d:Z/ejej01d;d<d#d=d$d>d?d@gdAdBdCdAdBdCgdddddd d!d"d d!g
ej2dDdEd=dFd@e3e4ddGd fdHdIZ5ej%&dJd$d>d?dKgej%&de5e5j6dAdLgdddMdNZ7ej%&de5e5j6dAdLgdddOdPZ8ej%9dQdddRdSZ:dddTdUZ;dddVdWZ<dddXdYZ=dddZd[Z>ddd\d]Z?ddd^d_Z@ddd`daZAdddbdcZBdddddeZCdddfdgZDdddhdiZEej%&djd6d5gdddkdlZFdddmdnZGej%&dod6d5gdddpdqZHdddrdsZIdddtduZJdddvdwZKdddxdyZLdddzd{ZMd|d} ZNd~d ZOdd ZPdd ZQdd ZReej%&dddiddigdd ZSG dd dZTG dd dZUG dd dZVddddZWddddZXej%&ddddddddgfdddddddgfgddddddZYdS )    )annotationsN)	to_offset)	DataArrayDatasetVariable)_consolidate_slices)assert_allcloseassert_array_equalassert_equalassert_identicalcreate_test_datarequires_daskrequires_floxrequires_scipyc                  C  sn   t dtjdddfddddgfd	dd
ddgddifddddgddgd} dddddggd f| d< | S )Nxyz         r   efg)fooZbazabcname   )r   r   hjboo)xrr   nprandomrandnds r)   =/tmp/pip-unpacked-wheel-h316xyqg/xarray/tests/test_groupby.pydataset   s    &r+   c                 C  s   | d S )Nr   r)   r+   r)   r)   r*   array'   s    r-   Nonereturnc               	   C  s   t tdtddgtdgks$tt tddtddgtddgksLtt tdddtdddgtdddgkszttddtddg} t | | ksttt t tddg W 5 Q R X d S )Nr      r      r   r   )r   sliceAssertionErrorpytestraises
ValueError)Zslicesr)   r)   r*   test_consolidate_slices,   s    $(.r8   c                 C  sl   |  dj| jddjkst|  dj| jddjks<t| ddi}| dj|jdd	jkshtd S )
Nr   r   r   r   r   xyr   r   r   )r;   )groupbydimsiselr4   stack)r+   stackedr)   r)   r*   test_groupby_dims_property9   s    rB   c                 C  sP   | j ddddg }d| }|jddgdd	d
d d	}t|| d S )Nr   T)r   dropr   r   r   r   spacerE   c                 S  s   d|  S Nr   r)   r9   r)   r)   r*   <lambda>H       z.test_multi_index_groupby_map.<locals>.<lambda>)r?   r@   r=   mapunstackr
   )r+   r(   expectedactualr)   r)   r*   test_multi_index_groupby_mapA   s    rM   c              	   C  s\   | j ddd}tjdd | }W 5 Q R X tjdd | }W 5 Q R X t|| d S )Nr   FsqueezeZuse_floxT)r=   r#   set_optionssumr   )r+   gbrK   rL   r)   r)   r*   test_reduce_numeric_onlyN   s    rT   c                  C  sh   t ddtdfidddgddd	d
gd} | d}| jddgdddd}t|| d S )Nr   r   )r   r   r   r   r   r   r   r   r   r   r<   r   r   r   rD   rE   )	r#   r   r$   onesrR   r@   r=   rJ   r
   )r(   rK   rL   r)   r)   r*   test_multi_index_groupby_sumW   s    
"rV   c            	      C  s   t jddd} tjddddgt| ddd	}|  }|d
 |d g}|d
d d |dd d  }tj|t| dddd}||}|jdd}tjddgt|ddd	}t|| d S )N
2000-01-01r   periodsr   r   r   timer[   coordsr>   r   reference_dater]   r>   r   dim   )r^   )	pd
date_ranger#   r   dictZto_pydatetimer=   rR   r
   )	timesr   ddZreference_dateslabelsindr   rL   rK   r)   r)   r*   test_groupby_da_datetimeb   s&        
  rj   c                  C  sT   t dddgddddgfg} t ddgdddgfg}| d }t|| d S Nr   r   r   r   )r#   r   r=   rR   r
   )r-   rK   rL   r)   r)   r*   (test_groupby_duplicate_coordinate_labelsv   s    rl   c                  C  sf   t dddgddddgfg} |  }t ddgdddgfg}| d }t|| t| | d S rk   )r#   r   copyr=   rR   r   )r-   Z
array_copyrK   rL   r)   r)   r*   test_groupby_input_mutation~   s    
rn   objr   r   r   r   r1   r2   r   r   c                 C  s6   | j ddddgd}| ddd }t|| d S )	Nr   r   r   r   r9   r   c                 S  s   | j ddgdS )Nr   r   r9   )r?   )r   r)   r)   r*   rG      rH   z0test_groupby_map_shrink_groups.<locals>.<lambda>r?   r=   rI   r   )ro   rK   rL   r)   r)   r*   test_groupby_map_shrink_groups   s    rq   c                 C  s8   dd }| j dddgd}| d|}t|| d S )Nc                 S  s2   | j d dkr | jddgd}n| jdgd}|S )Nr   r   r   r9   )Zsizesr?   )groupresultr)   r)   r*   func   s    z0test_groupby_map_change_group_size.<locals>.funcr   r   r9   r   rp   )ro   rt   rK   rL   r)   r)   r*   "test_groupby_map_change_group_size   s    ru   c                  C  sj   d
dd} t dddgddddgfg}t dddgddddgfg}|dj| ddd	}t|| d S )Nr   c                 S  s   | | | S Nr)   Zarg1Zarg2arg3r)   r)   r*   rt      s    z+test_da_groupby_map_func_args.<locals>.funcr   r   r   r   r   argsrx   )r   )r#   r   r=   rI   r   )rt   r-   rK   rL   r)   r)   r*   test_da_groupby_map_func_args   s
    
r|   c                  C  sv   ddd} t dddddgfiddddgi}t dddddgfiddddgi}|dj| d	dd
}t|| d S )Nr   c                 S  s   | | | S rv   r)   rw   r)   r)   r*   rt      s    z+test_ds_groupby_map_func_args.<locals>.funcr   r   r   r   r   ry   rz   )r   )r#   r   r=   rI   r   )rt   r+   rK   rL   r)   r)   r*   test_ds_groupby_map_func_args   s
    
$$r}   c               	   C  s2   t jg dd} tt | d W 5 Q R X d S )Nra   r>   )r#   r   r5   r6   r7   r=   )Zempty_arrayr)   r)   r*   test_da_groupby_empty   s    r   c                  C  s  t jddddddgdddddddgidd} t jddgddgd	d
dd}| dd	}t|| t jddgddggddgddgd
d
d}| dddg}t|| t jtjdddddgdddddddgidd} dD ]V}|dkrtjdgnddg}t j|ddgd	d
dd}| djd	|d}t|| qt jdddgdddgdddgdddgdddggdddddgdddgddd} | djddd}t jddgddgdd
dd}t|| | djddd}t jddgddgdddd}t|| | dd}	t jdddgdddggddgdddgdddd}
t|
|	 | dd}t jddgddgddgddgddggdddddgddgdddd}t|| tjd d!d"}ddg}t jt	t
d#d$||d%d%d&}||jjj}|jddd}t jd'd(d)d*d+d,d-d.d/d0d1d2gt
dddd3d4d}t|| |jdd5dd d }t jd'dgd(d6ggddgddgdd7d8d}t|| t jddddgdddddgidd} t jddgddgd	d
dd}| djd	d9d:}t|| d S );Nr   r   r   r   r1   r2   r   datar]   r>         ?r   quantiler   TFNF      @skipna                              r<   .r`   r   r   r   r   r   r   rW   m  rX     r   r   r[   r   r\                 O@      ^@     f@     @n@      s@     v@     z@     ~@      @     @     @monthr   r   r[   ?   r   r   r   r   r   lowermethod)r#   r   r=   r   r   r$   NaNrc   rd   reshapearanger[   dtr   )r-   rK   rL   r   r   actual_x
expected_xactual_y
expected_y	actual_xxexpected_xx	actual_yyexpected_yyrf   r   r   r   r)   r)   r*   test_da_groupby_quantile   s        

  *  
  

 


$  r   c                  C  sV  t jddddddddgfidddddddgid	} t jddddgfid
ddgdd	}| dd
}t|| t jddddgddggfiddgddgdd	}| dddg}t|| t jddtjdddddgfidddddddgid	} dD ]^}|dkrtjdgnddg}t jdd|fid
ddgdd	}| djd
|d}t|| qt jdddddgdddgdddgdddgdddggfidddddgdddgdd	} | djddd}t jddddgfiddgddd}t|| | d jddd}t jdd ddgfiddgdd!d}t|| | dd}	t jdddddgdddggfiddgdddgdd"d}
t|
|	 | d d}t jddddgddgddgddgddggfidddddgddgdd"d }t|| tj	d#d$d%}ddg}t jdd&t
td'd(fit||d&d}||jjj}|jddd}t jdd)d*d+d,d-d.d/d0d1d2d3d4d5gfitdddd6d}t|| |jdd7djtd dd8}t jdd9d*dgd+d:ggfiddgddgdd;d	}t|| t jddddddgfidddddgid	} t jddddgfid
ddgdd	}| djd
d<d=}t|| d S )>Nr   r   r   r   r   r   r1   r2   )	data_varsr]   r   )r   r   r   r   r   Fr   r   r<   r   r   r   r   r   r   r   r   r   r   .r`   r]   r   r   r   rW   r   rX   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   	transposerc   rd   r   r   re   r[   r   r   r?   r3   )r(   rK   rL   r   r   r   r   r   r   r   r   r   r   rf   r   r   r   r)   r)   r*   test_ds_groupby_quantile9  s      


 *
$
$

(




* r   
as_datasetFTc              
   C  s   t jddddgdddddgidd}| r4|jddn|}tjtd	d
 |jddd}W 5 Q R X |jddd}t|| tj	dd0 tj
tdd
 |jdddd W 5 Q R X W 5 Q R X d S )Nr   r   r   r   r   r   r   r   z<`interpolation` argument to quantile was renamed to `method`matchr   r   )interpolationr   T)recordz!interpolation and method keywords)r   r   )r#   r   
to_datasetr5   warnsFutureWarningr   r   warningscatch_warningsr6   	TypeError)r   r-   arrrL   rK   r)   r)   r*   .test_groupby_quantile_interpolation_deprecated  s    $
r   c                  C  s   t jdddgdddggddgtd	tdd
d} | ddddgi}| djddgd}t jdddgdddggddgddgtdd
d}t|| t|| d S )Nr   r   r1   r2   rb      r   r   r   )r   r   r>   r]   r:   )r#   r   ranger=   assign_coordsr   )rL   actual1actual2rK   r)   r)   r*   test_da_groupby_assign_coords  s        
r   
      r   r   r   tr   r   r   
2001-01-01MfreqrY   r   )r   r   r   r   r   ra   r   r   c                 C  s   t | |}| jj d}|d| 7 }|dtt| |  7 }|dkrT|d7 }n4|dkrf|d7 }n"|dkrx|d	7 }n|d
kr|d7 }||kstd S )NGroupByz, grouped over %r
%r groups with labels r   z1, 2, 3, 4, 5.r   z*0, 1, 2, 3, 4, 5, ..., 15, 16, 17, 18, 19.r   z'a', 'b', 'c'.r   &1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12.)reprr=   	__class____name__lenr$   uniquer4   )ro   ra   rL   rK   r)   r)   r*   test_groupby_repr  s    


r   c                 C  sX   t | d}| jj d}|d7 }|dtt| jjj	 7 }|d7 }||ksTt
d S )Nzt.monthr   z, grouped over 'month'r   r   )r   r=   r   r   r   r$   r   r   r   r   r4   )ro   rL   rK   r)   r)   r*   test_groupby_repr_datetime  s    r   z9ignore:invalid value encountered in divide:RuntimeWarningc               
   C  s  t jdtddfdtddfdtdtd	td
dd} tj| d jd< tj| d jd< tj| d jd< | | j}| 	 }tj|j
jddd d f< tj|j
jddd d f< tj|j
jd	dd d f< |dd j| j
j }t|| | }| dddgi}|j
|j ddi jdddj|jjddj|j }t|| |d}| d| j }t|| t ddd	gdddd	gfg}	dddtjgf|	d< t d	ddgfg}
|	d }t|
| dtdtdtd gf|	d!< t d	d!tdgfg}|	d! }t|| t ddddd	dgdtjddtjdtjgfg}	t d	d	gdddgfg}|	d }t|| d S )"N)latlonr[   g      N@)r   r   r1   )r   r   g      (@)r   r   )variableidr   r   r1   r   r   )r   r   )r   r   )r   r   r   r   c                 S  s   | S rv   r)   r9   r)   r)   r*   rG     rH   z)test_groupby_drops_nans.<locals>.<lambda>r;   r   r   T)rC   )r   r[   r   r   r   x1r   ZNaTr   )r#   r   r$   r   r   nanvaluesr=   r   rm   r   rI   r   r>   r   meanr@   whereZnotnullrenamer   Zreset_indexassignZdropnar   rR   r
   
datetime64)r(   grouped	expected1r   r   rA   	expected2Zactual3Z	expected3r-   Z	expected4Zactual4Z	expected5Zactual5Z	expected6Zactual6r)   r)   r*   test_groupby_drops_nans  sj    
 	




0r   c               	   C  sP  t dddddgfiddddgi} tjtdd | jddd	d
gd W 5 Q R X tjtdd |  jddd	d
gd W 5 Q R X tjtdd  | jdtjtjtjgd W 5 Q R X tjtdd$ |  jdtjtjtjgd W 5 Q R X tjtdd | 	| j
tj  W 5 Q R X tjtdd |  	| j
tj  W 5 Q R X d S )Nr   r   r   r   r   z-None of the data falls within bins with edgesr   g?g?g333333?binszAll bin edges are NaN.zFailed to group data.)r#   r   r5   r6   r7   groupby_binsZto_arrayr$   r   r=   r   r,   r)   r)   r*   test_groupby_grouping_errorsB  s&    $  "$(r   c              	   C  s   |  d}tjtdd |  W 5 Q R X tjtdd |d W 5 Q R X tjtdd |d W 5 Q R X | j ddd}t| |  t| d|tjd t| dd	g|tjdd	g d S )
Nr   zcannot reduce over dimensionsr   Zhuh)r   r   asdFrN   r   r   )	r=   r5   r6   r7   r   r   reducer$   r   )r-   r   r)   r)   r*   #test_groupby_reduce_dimension_error[  s    
r   c              	   C  s&   t t | dd W 5 Q R X d S )Nr   r   )r5   r6   r   r=   r-   r)   r)   r*   !test_groupby_multiple_string_argsm  s    r   c                  C  s   t  } t jtjdddddd| d< t jt| d jdd| d< tjdddd	}| d|	 }t jd
td dgdt
||jidjdd}t|| d S )Nz
2010-08-01z
2010-08-15Z15minr   r[   r~   val24H)startendr   `   )   	time_binsr   r   )r#   r   r   rc   rd   r$   rU   shaper   rR   cut
categoriesr   r   )r(   r  rL   rK   r)   r)   r*   test_groupby_bins_timeseriesr  s      
r
  c                  C  sJ   t dd } t| }tt | d}|| }d|jksFtd S )Nr   r   rr   )	r$   r   r#   r   Zfloor_divider=   r   r>   r4   )r   dakeyr   r)   r)   r*   test_groupby_none_group_name  s
    
r  c                 C  s   t | jdd| dd  t | jdd| dd  t | jjdd| jdd  t | jjdd| jdd  | dd  dd	d}| jdgdd
gddd	d}t || d S )Nr   r9   r   r   r   r   r"   r   r   r   )r   r   )r   selr=   r   rJ   r   )r+   rL   rK   r)   r)   r*   test_groupby_getitem  s      r  c            	      C  s  t dddgtjddfidtdfddddgftdd	} | d}t|dksXtddd
d}|j	|ksrtd| j
ddfd| j
ddfd| j
d
dfg}t||D ].\}}|d |d kstt|d |d  qdd }d	D ]"}| j|dd|}t| | qd S )Nr   r   r   r   r1   abcr   r   )r   r   r   r   r   r   r   r   r9   r   r   c                 S  s   | S rv   r)   r9   r)   r)   r*   identity  s    z&test_groupby_dataset.<locals>.identityFrN   )r   r$   r%   r&   listr   r=   r   r4   groupsr?   zipr
   rI   )	r   r=   expected_groupsZexpected_itemsr   r   r  kr   r)   r)   r*   test_groupby_dataset  s&     
r  c                  C  sn   t dddgtjddfi} | ddd }| d }t|| | d ddd }| }t|| d S )	Nr   r   r   r   r1   c                 S  s   | d S )Nr   r)   r'   r)   r)   r*   rG     rH   z7test_groupby_dataset_returns_new_type.<locals>.<lambda>c                 S  s   |   S rv   )r   r9   r)   r)   r*   rG     rH   )r   r$   r%   r&   r=   rI   r   )r   r   r   r   r   r)   r)   r*   %test_groupby_dataset_returns_new_type  s    
r  c                  C  s   t  } tt| dd d D ]j\}\}}| d | |ks@tt| d | |d  t| d | |d  t| d d d |f |d  q d S )Ndim1r   Zvar1Zvar2Zvar3)r   	enumerater  r=   r4   r
   )r   nr   subr)   r)   r*   test_groupby_dataset_iter  s    &r  c               	   C  s   t  } tjtdd | td W 5 Q R X tjtdd | | d d d  W 5 Q R X tjtdd | | jd 	  W 5 Q R X d S )Nz`group` must ber   r   zlength does not matchr  r   )
r   r5   r6   r   r=   r$   r   r7   r]   to_index)r   r)   r)   r*   test_groupby_dataset_errors  s     r!  c               	   C  s  t ddgtjddfdtjdfdtjdfdddddgfd} | d}|d jddi|d< | dd	}t|| | dd}t|| | d
 }t | d |d	| d  jd
di| d | d}| d
d	}t|| d S )Nr   r   r   r   r   r   )r;   xonlyyonlylettersr#  .r$  r;   r"  r   )r;   r"  r#  )	r   r$   r%   r&   r   r   Zset_dimsr=   r   )r   rK   rL   r$  r)   r)   r*   test_groupby_dataset_reduce  s,    	


r%  rO   c                 C  s   dd }t  }|d |d< |jd| d}|||jd  }||jd  }t||| |jd | }t||| d| }||| }|| }t||| || }t||| d S )Nc                 S  s   |  ddddS )Nr  dim2dim3r[   )r   r9   r)   r)   r*   reorder_dims  s    z/test_groupby_dataset_math.<locals>.reorder_dimsr  rN   r   )r   r=   r]   r   )rO   r(  r(   r   rK   rL   Zds2r)   r)   r*   test_groupby_dataset_math  s    r)  c               	   C  sh  t  } | d}tddddgdtdfg}| tdtd dddd}|| }t|| || }t|| t	j
td	d
 ||   W 5 Q R X t	j
td	d
 | |  W 5 Q R X t	j
tdd
 |d  W 5 Q R X t	j
tdd
 ||  W 5 Q R X t	j
tdd
 | |7 } W 5 Q R X tdtdftjdddd} t	j
td	d
 | | d  W 5 Q R X d S )NZnumbersr   r   r'  r   r  r&  r[   zincompat.* grouped binaryr   only support binary opsr   in-place operationsd   rW   rX   )r   r[   z
time.month)r   r=   r   r   r   r$   zerosr   r
   r5   r6   r7   r   r   r   rc   rd   )r(   r   r-  rK   rL   r)   r)   r*   test_groupby_math_more  s<    
   

r.  indexed_coordc              	   C  s   d}t tj||fdd}| r<t||d< t||d< |dtd|d d}| }|jtdd d	|jdddddddgfd
 }|| }t|| d S )Nrb   r<   r~   r   r   r   r   r   r9   )x_bins)	r   r$   r%   r   r   r   r?   r3   r   )r/  Nr  r   r   rK   rL   r)   r)   r*   test_groupby_bins_math-  s    .r2  c                  C  s  d} t tj| | fdddtjddddd	d
ddg| d dfid}t|j|d |d< |d}| }||j	|j
d }|j|j
|d< || }t|| dtjddddddddg| d df|d< t|j|d |d< |jddddgd}| }tj|jdddd d d d f d }|d!jd|fd"}|jtd d#| }|j|j
|d< |j|j|d< d|jj| f|d!< || }t|| d S )$N(   r<   rh   r   r   r   r   dr   r   r   r    r   )Zrepeatsr   r   labels2d)r5  r   r   r   r   r1   r2   rb   numnum2dr   )r   r   r2   T)r   right   
num2d_bins)r:  r9   )r   r$   r%   repeatr#   	broadcastrh   r=   r   r  r5  Zbroadcast_liker   r6  r   Zdigitizer7  	drop_varsr?   r3   r:  r   )r1  r  r   r   rK   rL   ZidxrZexpanded_meanr)   r)   r*   test_groupby_math_nD_group;  s>     


 &r>  c                  C  sl   t dddddgfidtjdddi} | d}||d	 }t ddd
d
d
gfi| ddg }t|| d S )Nr   r   r   r   r   Z20100101rX   zt.day.r   )r   rc   rd   r=   r   r   )r(   r   rL   rK   r)   r)   r*   !test_groupby_dataset_math_virtuala  s
    &
 r?  c               	   C  sN   t tdddtjddddid} | d	}||  }|j| jksJtd S )
N)r   r   r   r   r   r[   r[   r   r   6HrY   r   r   ztime.day)	r   r$   rU   rc   rd   r=   r   r>   r4   )r  r   rs   r)   r)   r*   test_groupby_math_dim_orderi  s    
rC  c                  C  sd   t ddddddgfidddddtjgfi} | dd}t dd	dgfddgd
}t|| d S )Nr   r   r   r   r   r   bar.      ?)r   rD  )r   r$   r   r=   r   r   r(   rL   rK   r)   r)   r*   test_groupby_dataset_nant  s    ,rG  c                  C  sd   t  } dD ]}ttddgd| |< q
t| j }| dd} t| j }||ks`t	d S )Nr  r   r   r~   .)
r   r   r$   r   r  r   keysr=   r   r4   )r(   ZvnZdata_vars_refr   r)   r)   r*   test_groupby_dataset_order|  s    rI  c                  C  s4  t ddtjdtjdgfidddddgi} t ddtdfidddddgi}| |fD ]}dddddgf|jd< qT| dtddgdd	}t|| | dt ddddgfi}t|| d
| j	d< d| j
j	d< | dt ddddgfi}|j	| j	kst|j
jdkst|j
j	| j
j	ks0td S )Nr   r   r   r   r   r   r   r   r~   r(   attrr  )r   r$   r   r   r]   r=   fillnar   r   attrsr   r4   r   )r(   rK   targetrL   r)   r)   r*   test_groupby_dataset_fillna  s    *"
 

 rN  c               	   C  s   t ddtdfidddddddgfi} t ddddgfi}| jdd	}ddgtjgd
  |d _| d|}t|| d| j	d< d| j
j	d< | d|}|j	| j	kst|j
jdkst|j
j	| j
j	kstd S )Nr   r   r1   r   r   r   TF)deepr   r(   rJ  r  )r   r   rm   r$   r   r   r=   r   r   rL  r   r4   r   )r(   ZcondrK   rL   r)   r)   r*   test_groupby_dataset_where  s    (

rP  c                  C  s   t ddtdfidddgd dg fi} | djdd	 d
}| dddddgfi}t|| | djdd	 d
}| dddddgfi}t|| | djdd	 d
}|d}t|| d S )Nr   r   r   r   Ar   Bc                 S  s
   d| j  S rF   r   r'   r)   r)   r*   rG     rH   z-test_groupby_dataset_assign.<locals>.<lambda>)r   r   r   r   c                 S  s
   | j  S rv   r   rR   r'   r)   r)   r*   rG     rH   r   c                 S  s
   | j  S rv   rT  r'   r)   r)   r*   rG     rH   )r   r   r=   r   merger   r   
set_coordsrF  r)   r)   r*   test_groupby_dataset_assign  s    *


rW  c                  C  sf   t ddddddgfidddddgid} | dd	d
 }tddgdddgiddd}t|| d S )Nr   r   r   r   r   r   r   r   c                 S  s
   | j  S rv   )r   r   grpr)   r)   r*   rG     rH   z9test_groupby_dataset_map_dataarray_func.<locals>.<lambda>rE        @r_   )r   r=   rI   r   r   rF  r)   r)   r*   'test_groupby_dataset_map_dataarray_func  s    (r[  c                  C  sh   t ddddgdddddgiddd} | dd	d
 }tjddddgfidddgid}t|| d S )Nr   r   r   r   r   r   r   r_   c                 S  s   |    S rv   )r   r   rX  r)   r)   r*   rG     rH   z9test_groupby_dataarray_map_dataset_func.<locals>.<lambda>rE  rZ  r   )r   r=   rI   r#   r   r   )r  rL   rK   r)   r)   r*   'test_groupby_dataarray_map_dataset_func  s    $"r\  kwargsr   z
map-reduceZenginenumpyc              	   C  s   t ddtdfidddddddgfi}tjdd |d }W 5 Q R X tjd	d |djf | }W 5 Q R X t|| d S )
Nr   r   r1   r   r   r   FrP   T)r   r   r#   rQ   r=   r   r   )r]  r(   rK   rL   r)   r)   r*   test_groupby_flox_kwargs  s    (r_  c                   @  s  e Zd Zejdddd Zdd Zdd Zd	d
 Zej	
dddddgej	
dddgej	
dddgddddZdd Zej	
ddddgdd Zdd  Zej	
dddgej	
d!d"ddgd#d$ Zej	
d!d"ddgd%d& Zd'd( Zd)d* Zd+d, Zej	
dddgd-d. Zd/d0 Zd1d2 Zd3d4 Zd5d6 Zd7d8 Zd9d: Zd;d< Zd=d> Zd?d@ ZdAdB ZdCdD ZdEdF Z dGdH Z!dIdJ Z"d"S )KTestDataArrayGroupByT)Zautousec                 C  s   ddd| _ tjd| _tddg| j| _tddg| j| j | _td| ji| _| jd | _	t
jjdd	gd
dggdd| _tdd
ddgd| jidd| _| j	 | _dtdgd dg d	gd  f| jjd< dd| jd   | jjd< d S )NZvalue1iq  )Zattr1Zattr2)r   r   r   r   r   r   r   r   r   )Zlevel_1Zlevel_2namesr   r   r\   	   r   r   r  r   r,  )rL  r$   r%   r   r   vvar   r(   dvrc   
MultiIndexfrom_productZmindexr   Zmdarm   r  r-   r]   selfr)   r)   r*   setup  s     ,zTestDataArrayGroupBy.setupc                 C  s  ddgddgg}ddddg}ddg}ddg}t j||d|id}|j|dd	 }tjjddgddgg|d
}t j|d	gd	|id}t|| t j||d|d d d id}|j|dd	 }	tjjddgddgg|d
}
t j|d	gd	|
id}t|	| d S )Nr   r   r   r   r   r   r   r  r   ra  r   )	r#   r   r@   r=   firstrc   rg  rh  r
   )rj  r   Z	data_flatr>   Zy_valsr   r   Zmidx1r   r   Zmidx2r   r)   r)   r*   !test_stack_groupby_unsorted_coord  s    
z6TestDataArrayGroupBy.test_stack_groupby_unsorted_coordc                 C  sv   t | jd| jdD ]*\\}}\}}||ks6tt|d | qt | jd| jD ]\\}}}t|| qZd S )Nr   r   r   )r  rf  r=   r(   r4   r   )rj  Zact_xZact_dvZexp_xZexp_ds_Zexp_dvr)   r)   r*   test_groupby_iter  s    
 
 z&TestDataArrayGroupBy.test_groupby_iterc                 C  sn   | j d}tdddgtddd}| |j ks<t|D ]}t|| |j|  q@dt|ksjtd S )Nr  r   rc  r   r   )r   r   r   r   )r  r=   r   rH  r  r4   r	   r   )rj  r   r  r  r)   r)   r*   test_groupby_properties  s    z,TestDataArrayGroupBy.test_groupby_propertiesz
by, use_da)r   F)r   F)r   T)r  FshortcutFrO   r.   r/   c           	      C  sF   | j }|r|j| }dd }|j||d}|j||d}t|| d S )Nc                 S  s   | S rv   r)   r9   r)   r)   r*   r    s    z@TestDataArrayGroupBy.test_groupby_map_identity.<locals>.identityrN   )rq  )r  r]   r=   rI   r   )	rj  byZuse_darq  rO   rK   r  r   rL   r)   r)   r*   test_groupby_map_identity  s    
z.TestDataArrayGroupBy.test_groupby_map_identityc                   s  | j   d}ttdgt| jd d d df  | jd d dd f  | jd d ddf  gjtdgtdddgdd }t	||j
tjd	d
 t	||d	 t fddtdtdd tddfD dddggdg} d dtj}t	||  d dd	}t	|| tddgt| jd d d df d| jd d dd f d| jd d ddf dgjftdgtdddgdd }t	||
tjd t	||d d S )Nr  rc  r   r   r   r   r   r  r   .r`   c                   s   g | ]} d  j |  qS r:   )r   rR   ).0idxr   r)   r*   
<listcomp>;  s   z9TestDataArrayGroupBy.test_groupby_sum.<locals>.<listcomp>r   r   r   )r  r=   r   r   r$   r-   r   rR   Tr   r   r   r3   rI   )rj  r   Zexpected_sum_allrK   rL   Zexpected_sum_axis1r)   r   r*   test_groupby_sum$  s\    






z%TestDataArrayGroupBy.test_groupby_sumr   rR   r   Zmedianc                 C  s  | j }|d}tt|}ttddgt|| jd d d df dd|| jd d dd f dd|| jd d ddf ddgjtdgtddd	gd
d }t	j
dd t||dd}W 5 Q R X t	j
dd t||dd}W 5 Q R X t|| t|| d S )Nr  r   rc  r   axisr   r   r   r   rt  r   FrP   r   r`   T)r  r=   getattrr$   r   r   r-   r   rx  r#   rQ   r   )rj  r   r-   r   Z	reductionrK   Zactual_legacyZ
actual_npgr)   r)   r*   test_groupby_reductionsY  s0    



z,TestDataArrayGroupBy.test_groupby_reductionsc              
   C  sl   t ddtjtjddgddddddddgfidd}|d }t ddd	gddddgfgd
}t|| d S )Nr   catr   r   r   r   r\   r   r   r   )r   r$   r   r=   countr   rj  r-   rL   rK   r)   r)   r*   test_groupby_countx  s    z'TestDataArrayGroupBy.test_groupby_count
keep_attrsNc              	   C  s`   | j }d|jd< |djtj||d}tjdd |dj|d}W 5 Q R X t|| d S )NrD  r   r  )r  rq  FrP   r  )	r  rL  r=   r   r$   r   r#   rQ   r   )rj  rq  r  r-   rL   rK   r)   r)   r*   test_groupby_reduce_keep_attrs  s    

  z3TestDataArrayGroupBy.test_groupby_reduce_keep_attrsc              	   C  sz   | j }d|jd< tjdd |dj|d}W 5 Q R X tjdd |dj|d}W 5 Q R X |j|_t|| d S )NrD  r   FrP   r  r  T)r  rL  r#   rQ   r=   r   r   r   )rj  r  r-   rK   rL   r)   r)   r*   test_groupby_keep_attrs  s    
z,TestDataArrayGroupBy.test_groupby_keep_attrsc              	   C  s   dd }| j }|d}| }t|| jd d d df || jd d ddf || jd d dd f g}ddg|f|d< |d }t||| d S )	Nc                 S  s   | t |  S rv   )r$   r   r9   r)   r)   r*   center  s    z<TestDataArrayGroupBy.test_groupby_map_center.<locals>.centerr  rc  r   r   r   r   )r  r=   r   r$   Zhstackr   r   rI   )rj  r  r-   r   Zexpected_dsZexp_dataZexpected_centeredr)   r)   r*   test_groupby_map_center  s    
Jz,TestDataArrayGroupBy.test_groupby_map_centerc                 C  s*   | j }|d}|tj}t|| d S )Nr  )r  r=   rI   r$   Zasarrayr
   )rj  r-   r   rL   r)   r)   r*   test_groupby_map_ndarray  s    
z-TestDataArrayGroupBy.test_groupby_map_ndarrayc                 C  s@   dd }| j }|d}||}| }||}t|| d S )Nc                 S  s"   | j d d | j d< d| jd< | S )Nr   r   ZlemonZfruit)r]   rL  r9   r)   r)   r*   change_metadata  s    
zOTestDataArrayGroupBy.test_groupby_map_changes_metadata.<locals>.change_metadatar  )r  r=   rI   rm   r
   )rj  r  r-   r   rL   rK   r)   r)   r*   !test_groupby_map_changes_metadata  s    

z6TestDataArrayGroupBy.test_groupby_map_changes_metadatac                 C  s   | j }|jd|d}||jd  }||jd  }t|| |jd | }t|| |jd jdd}|| }|| }t|| || }t|| d S )Nr   rN   Xr   )r  r=   r]   r   r   )rj  rO   r-   r   rK   rL   r(   r)   r)   r*   test_groupby_math_squeeze  s    


z.TestDataArrayGroupBy.test_groupby_math_squeezec              	   C  s   | j }|d}|dtd d }|ttdddddgfg }|dd}t|| t	j
tdd |d	  W 5 Q R X t	j
tdd ||  W 5 Q R X t	j
td
d ||7 }W 5 Q R X d S )Nr  .r   r   r   r   r*  r   r   r+  )r  r=   r   r$   r   r   r   r   r   r5   r6   r   )rj  r-   r   Zexpected_aggrL   Z
actual_aggr)   r)   r*   test_groupby_math  s    

z&TestDataArrayGroupBy.test_groupby_mathc                 C  s  t tddddddgfddddgddd}t d	gd
dgid
d}|d
| }t d	dtjtjg|j}t|| t d	gddgdd
d}|d
| }dgdgd tjgd  f|jd< t|| tdd
d	gfid
dgi}|d
| }tddd	dtjtjgfi|j}t|| d S )Nr   r   r   r   r   r   )r   r   r~   r   r   r\   r   {   )r   r   r   r   )r   r   r=   r$   r   r]   r   r   )rj  r-   otherrL   rK   r)   r)   r*   test_groupby_math_not_aligned  s"      
"
 z2TestDataArrayGroupBy.test_groupby_math_not_alignedc                 C  sf   t tjdddtdfdtdfdddgd}dD ]*\}}||dd	 }|j|ks6tq6d S )
Nr1   r   r   r   )r   r   r\   )r   r<   )r   r<   )r   )r   r   )r   )r   r   c                 S  s   |   S rv   rN   r9   r)   r)   r*   rG     rH   zETestDataArrayGroupBy.test_groupby_restore_dim_order.<locals>.<lambda>	r   r$   r%   r&   r   r=   rI   r>   r4   rj  r-   rr  Zexpected_dimsrs   r)   r)   r*   test_groupby_restore_dim_order  s    z3TestDataArrayGroupBy.test_groupby_restore_dim_orderc              	   C  s~   t tjdddtdfdtdfdtjddfdddgd}dD ]2\}}|j|d	d
dd d }|j|ksFtqFd S )Nr1   r   r   r   r<   r  r\   r  T)restore_coord_dimsc                 S  s   |   S rv   rN   r9   r)   r)   r*   rG     rH   zFTestDataArrayGroupBy.test_groupby_restore_coord_dims.<locals>.<lambda>r   r  r  r)   r)   r*   test_groupby_restore_coord_dims  s    


z4TestDataArrayGroupBy.test_groupby_restore_coord_dimsc                 C  s  t dddddgdd}t dgd d	gd  dd
d}t ddgd
dd	gfg}|| }t|| t ddgd
dd	gfg}|| }t|| t tjddddgd}t |ddg d
dd	gid
dg}|| }t|| |d }|}t|| d S )Nr   r   r   r   r1   r   r~   r   r   ab)r>   r   r   r   )r   r=   rl  r   lastr$   r%   r&   )rj  r-   rr  rK   rL   r)   r)   r*   test_groupby_first_and_last  s    

 
z0TestDataArrayGroupBy.test_groupby_first_and_lastc                 C  s^   t ddgddggddgddgggd	d
gddgddggfd	d
gddgddggfddd	d
gdS )Nr   r   r   r   r1   r   r   r   nyZnxr9  r3  2   )r   r   r[   r\   )r   ri  r)   r)   r*   #make_groupby_multidim_example_array1  s    z8TestDataArrayGroupBy.make_groupby_multidim_example_arrayc                 C  sp   |   }dtdddgddddgfgdfd	td
dgd	ddgfgdffD ]"\}}||d}t|| qHd S )Nr   r1      r   g      >@g      D@g      I@r   r   r   r3        $@      4@.)r  r   r=   rR   r   )rj  r-   ra   Zexpected_sumZ
actual_sumr)   r)   r*   test_groupby_multidim;  s     z*TestDataArrayGroupBy.test_groupby_multidimc                 C  sX   |   }|ddd }tddgddggdd	gd
dggg|j|jd}t|| d S )Nr   c                 S  s   | |    S rv   r   r9   r)   r)   r*   rG   F  rH   z@TestDataArrayGroupBy.test_groupby_multidim_map.<locals>.<lambda>g      g      g      g      !r         @g       @g      !@r\   )r  r=   rI   r   r]   r>   r   r  r)   r)   r*   test_groupby_multidim_mapD  s    z.TestDataArrayGroupBy.test_groupby_multidim_mapc              	   C  sH  t tddd}d|d< dddg}t|d |j}t ddgd	d	|id
}|jd|d }t|| |jd|ddgd }t|j	ddgd| |jd|d
dd }t|| t|jdkstt td}ddddg}tjdd |d|d}W 5 Q R X tjdd |d|d}W 5 Q R X t|| d S )Nr   dim_0r~   c   r   rE  r1   r   
dim_0_binsr   r   g333333?rZ  )r   rh   )r  c                 S  s   |   S rv   rR   r9   r)   r)   r*   rG   _  rH   z8TestDataArrayGroupBy.test_groupby_bins.<locals>.<lambda>)r   r   r   r   r   FrP   .T)r   r$   r   rc   r  r	  r   rR   r   r   rI   r   r  r4   r#   rU   rQ   r   r   )rj  r-   r   
bin_coordsrK   rL   r  r)   r)   r*   test_groupby_binsN  s.    
  

z&TestDataArrayGroupBy.test_groupby_binsc                 C  s~   t tddtdfg}dddg}t|d |j}|d| }t dtj	gdd|id}t
|| t|jdksztd S )Nr   r   r   r1   r2   r0  r   )r   r$   r   r   rc   r  r	  r   rR   r   r   r   r   r4   )rj  r-   r   r  rL   rK   r)   r)   r*   test_groupby_bins_emptym  s    

z,TestDataArrayGroupBy.test_groupby_bins_emptyc           	      C  sj  |   }dddg}t|d jj|j}tddgdd|id}|d|d	d
 }t	|| t
ddgddgg|d _tddgdd|id}|d|dd
 }t	|| dddddg}tt
ddd}tt
dddddgd gdddd}|j||d }t
jdd t|d d |dd  D td}tt
dt
jddgd d |id}t	|| d S )!Nr   r   r   r   r   r3  Zlat_binsr   c                 S  s   |   S rv   r  r9   r)   r)   r*   rG   ~  rH   zATestDataArrayGroupBy.test_groupby_bins_multidim.<locals>.<lambda>r  r  r  c                 S  s   |   S rv   r  r9   r)   r)   r*   rG     rH   r   r   r   r   )r1   r   r<   r~   g      r   rE  r   r1   r   c                 S  s    g | ]\}}t j||d dqS )r8  )closed)rc   ZInterval)ru  leftr8  r)   r)   r*   rw    s   zCTestDataArrayGroupBy.test_groupby_bins_multidim.<locals>.<listcomp>Zdtyper2   Z
group_bins)r  rc   r  r   Zflatr	  r   r   rI   r   r$   r-   r   rU   r   r  r  objectr   )	rj  r-   r   r  rK   rL   fieldrr  Zbincoordr)   r)   r*   test_groupby_bins_multidimy  s:    


  z/TestDataArrayGroupBy.test_groupby_bins_multidimc              	   C  s   t jtdddtjddddid}|jddd }| jsFt	t j
dd	 |jddd }W 5 Q R X t j
d
d	 |jddd }W 5 Q R X t|| d S )Nr,  r   i)r6  r   r   r   TrP   F)r#   r   r$   r   linspacer   r   r   Zis_monotonic_increasingr4   rQ   r  r   )rj  r   Zbinned_meanrL   rK   r)   r)   r*   test_groupby_bins_sort  s      z+TestDataArrayGroupBy.test_groupby_bins_sortc                 C  sh   t ddddgddddddgfidd}|djd	d
 d}| }dddddgf|jd< t|| d S )Nr   r   r   r   r   r   r   r~   c                 S  s   |   S rv   r  rS  r)   r)   r*   rG     rH   zATestDataArrayGroupBy.test_groupby_assign_coords.<locals>.<lambda>)r4  rE  rZ  r4  )r   r=   r   rm   r]   r   r  r)   r)   r*   test_groupby_assign_coords  s
    &z/TestDataArrayGroupBy.test_groupby_assign_coordsc                 C  s   t tjdtjdgdtdidd}t ddgdd}||}t ddgddgddgddggdtdid	d}t|| t tddtdidd}| }||fD ]}dddddgf|jd
< q|d
t ddgd
d}t|| d S )Nr   r   r   r   r\   r   r   r~   r<   r   r   )	r   r$   r   r   rK  r   rm   r]   r=   )rj  r   Z
fill_valuerL   rK   r   rM  r)   r)   r*   test_groupby_fillna  s    "
 
 
z(TestDataArrayGroupBy.test_groupby_fillna)#r   
__module____qualname__r5   fixturerk  rm  ro  rp  markparametrizers  ry  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`    sP   

	 
5




	
"r`  c                   @  s   e Zd Zdd Zdd Zdd Zdd Zed	d
 Zdd Z	dd Z
dd Zdd Zdd Zedd Zedd Zedd Zeeejdddgdd Zd d!d"d#Zd d!d$d%Zd d!d&d'Zd(S ))TestDataArrayResamplec              	   C  s   t jdddd}ttdd|fg}|jdd }t| d }t|| |jdd	tj}t|| |jddd	 }| d }| j
td7  _
t|}t|| tjtd
d |dddg jdd W 5 Q R X d S )NrW   rA  r   r   r[   r  rZ   -12Hr[   Zloffsetzindex must be monotonicr   r   r   r   1D)rc   rd   r   r$   r   resampler   	to_seriesr   r   indexr   from_seriesr5   r6   r7   )rj  rf   r-   rL   rK   	expected_r)   r)   r*   test_resample  s    



z#TestDataArrayResample.test_resamplec                 C  st   ddd}t jdddd}tjdddg|gd	gd
}tjdddg|gd	gd
}|jddj|ddd}t|| d S )Nr   c                 S  s   |  d| | S Nr[   r  rw   r)   r)   r*   rt     s    z>TestDataArrayResample.test_da_resample_func_args.<locals>.func2000r   DrB        ?r[   r\   r  rZ   r  rz   )r   )rc   rd   r#   r   r  rI   r   )rj  rt   rf   r  rK   rL   r)   r)   r*   test_da_resample_func_args  s    
z0TestDataArrayResample.test_da_resample_func_argsc                 C  s~  t jdddd}ttdd|fg}|jdd }tdd	d
gd|d d d	 fg}t|| |jdd }t| d }t|| |	t
}tj|d d< |jdd }tdd	d
gd|d d d	 fg}t|| |jddjdd}ttjd	d
gd|d d d	 fg}t|| td|id }|jdd }t dddg}t|d|d d d	 fgdd}t|| d S )NrW   rA  r   r   r[   r  rZ   r   r   r   r  r   Fr   z2000-01-01T18z2000-01-02T18z2000-01-03T06r   )rc   rd   r   r$   r   r  rl  r   r  astypefloatr   r   r  Zto_datetime)rj  rf   r-   rL   rK   expected_timesr)   r)   r*   test_resample_first  s.     


 
"
z)TestDataArrayResample.test_resample_firstc              	   C  sX   t jdddd}ttdd|fg}tjtdd |jf ddi	  W 5 Q R X d S )	NrW   rA  r   r   Z__resample_dim__zProxy resampling dimensionr   r  )
rc   rd   r   r$   r   r5   r6   r7   r  rl  )rj  rf   r-   r)   r)   r*   test_resample_bad_resample_dim  s    z4TestDataArrayResample.test_resample_bad_resample_dimc                 C  sB  t d}t d}tjdddd}t t dd}t |d |d \}}t jt|td	}t||||d
d}t|j	||dd}	t|j	||dd}
t|d|id}t
||	|
|d}|dddg}|d }|jdddd}d|jkst|jddd }d|jkst|jdddd}d|jks>td S )Nr2   r   rW   rA  r1   r   r2   r   r   r   r  r[   r   r   r@  r<   r[   rZ   r   xcyctcr  r  r  r   12HT)r[   r  1Hlinearr$   r   rc   rd   tileZmeshgridr   intr   rx  r   rV  r  r   r]   r4   ffillinterpolaterj  xsysrf   r   Zxxyyttr-   ZxcoordZycoordZtcoordr(   rL   r)   r)   r*    test_resample_drop_nondim_coords	  s*    

z6TestDataArrayResample.test_resample_drop_nondim_coordsc              	   C  s   t jdddd}ttdd|fg}d|jd< |jdd	jd
d}tdddgd|d d d fg|jd}t|| t	j
tdd |jdd
d W 5 Q R X d S )NrW   rA  r   r   r[   r   metar  rZ   Tr  r   r   )rL  5Passing ``keep_attrs`` to ``resample`` has no effect.r   r[   r  )rc   rd   r   r$   rU   rL  r  r   r   r5   r   UserWarningrj  rf   r-   rs   rK   r)   r)   r*   test_resample_keep_attrs)  s    
&
 z.TestDataArrayResample.test_resample_keep_attrsc                 C  st   t jdddd}ttdd|fg}tj|d< |jddjd	d
}ttjddgd|d d d fg}t|| d S )NrW   rA  r   r   r[   r   r  rZ   Fr   r   )	rc   rd   r   r$   rU   r   r  r   r   r  r)   r)   r*   test_resample_skipna7  s    
"z*TestDataArrayResample.test_resample_skipnac                 C  s   t jdddd}ttdd|fg}|jdd }t| d }t|| |jdd	 }t| d	 }t|| |jdd
 }t| d
 }t|| |jdd }t| d }t|| |jdd}| }|j}t|j|dd	}t|| d S )
NrW   rA  r1   r   r[   3HrZ   nearest)r[   r   )rc   rd   r   r$   r   r  r  r  r   bfillasfreqpadr  Z_full_indexreindex)rj  rf   r-   rL   rK   rsZ	new_timesr)   r)   r*   test_upsample@  s&    



z#TestDataArrayResample.test_upsamplec           
      C  s,  t d}t d}tjdddd}t t dd}t||||dd	}|jd
d }t j|ddd}|	 d

 j}|dd t|f }t||||dd	}	t|	| |jd
d }t jt |jjddd}t |jj}|	 d

 j}|dd t|f }t||||dd	}	t|	| |jd
d
 }t j|dddtdd df }t j|ddd df< |	 d

 j}t||||dd	}	t|	| |jd
d }t j|ddd}|dd d df |ddd df< |dd df }|	 d

 j}t||||dd	}	t|	| d S )Nr2   r   rW   rA  r1   r   r  r  r@  r  rZ   r   r   rz  .r   )r$   r   rc   rd   r  r   r  r  r;  r  r  r  r   r   Zflipudrx  r  r  r   r  )
rj  r  r  rf   r   r-   rL   expected_datar  rK   r)   r)   r*   test_upsample_nd_  s^    





"

 
z&TestDataArrayResample.test_upsample_ndc                 C  s   t jdddd}t jdddd}ttdd|fg}|jddjd	d
}tdddtjdgd|fg}t|| |jddj	d	d
}tdtjdddgd|fg}t|| |jddj
dd
}tddtjddgd|fg}t|| d S )NrW   r  r   r   rA  r1   r[   rZ   r  )Z	tolerancer   r  r   r   )rc   rd   r   r$   r   r  r  r   r   r  r  )rj  rf   Ztimes_upsampledr-   rL   rK   r)   r)   r*   test_upsample_tolerance  s    

z-TestDataArrayResample.test_upsample_tolerancec              	   C  s   ddl m} td}td}tjdddd}tdd	 }t|d
}t||||dd}| 	d
 j}tdt|d t|d }	dD ]^}
|j	dd|
}|tt|||
dddd}||	}t||||dd}t||dd qd S )Nr   interp1dr2   r   rW   rA  r1   r   r   r  r  r@  r  r   r  r  ZzeroZslinearZ	quadraticZcubicrZ   r   Tkindr{  Zbounds_errorZassume_sorted缉ؗҜ<Zrtol)scipy.interpolater  r$   r   rc   rd   r  r   r  r  r  r  r  r   r  r   )rj  r  r  r  rf   r   r   r-   r  new_times_idxr  rL   r   r  rK   r)   r)   r*   test_upsample_interpolate  s4    


z/TestDataArrayResample.test_upsample_interpolatec                 C  s~   t jdddd}ttt|d|fg}|jddd}t	t
d	t
d
g}tdtjgd|fg}t|| d S )Nz
2007-02-01z
2007-03-01r  r   r[   r   rZ   r  z
2007-02-28z
2007-03-31g      ;@)rc   rd   r#   r   r$   r   r   r  r  r-   r   r   r
   )rj  datesr  rs   r  rK   r)   r)   r*   "test_upsample_interpolate_bug_2197  s    z8TestDataArrayResample.test_upsample_interpolate_bug_2197c                 C  sZ   t jdddd}tjtjt|ddfdd|id	}|jdd
d}t	||dd d S )Nz
2016-01-01z
2016-03-31r  r   r   r   r  r[   r   rZ   r  r  r  )
rc   rd   r#   r   r$   r%   r   r  r  r   )rj  r  rK   rL   r)   r)   r*   )test_upsample_interpolate_regression_1605  s    z?TestDataArrayResample.test_upsample_interpolate_regression_1605chunked_timeTFc              	   C  s  ddl m} td}td}tjdddd}tdd	 }t|d
}t||||dd}d	dd}	|rtd|	d< | 	d
 j}
tdt|d t|d }dD ]l}||	j	dd|}| }|tt|||dddd}||}t||
||dd}t||dd qd S )Nr   r  r2   r   rW   rA  r1   r   r   r  r  r@  r   r<   r[   r  r  rZ   r   Tr  r  r  )r  r  r$   r   rc   rd   r  r   r  r  r  r  r  r   chunkr  Zcomputer   )rj  r  r  r  r  rf   r   r   r-   chunksr  r  r  rL   r   r  rK   r)   r)   r*   test_upsample_interpolate_dask  s<    



z4TestDataArrayResample.test_upsample_interpolate_daskr.   r/   c                 C  sd   t jdddd}ttdd|fg}d}|jd|d }t| jd|d	 }t|| d S )
N2000-01-01T02:03:01rA  r   r   r[   r   r  )r[   base)r  	rc   rd   r   r$   r   r  r   r  r   )rj  rf   r-   r  rL   rK   r)   r)   r*   test_resample_base  s    z(TestDataArrayResample.test_resample_basec                 C  sj   t jdddd}ttdd|fg}t d}|jd|d }t| jd|d	 }t	|| d S )
Nr  rA  r   r   r[   Z11Hr  )r[   offset)r
  )
rc   rd   r   r$   r   Z	Timedeltar  r   r  r   )rj  rf   r-   r
  rL   rK   r)   r)   r*   test_resample_offset  s    
z*TestDataArrayResample.test_resample_offsetc                 C  sd   t jdddd}ttdd|fg}d}|jd|d }t| jd|d	 }t|| d S )
Nr  rA  r   r   r[   r  r  )r[   origin)r  r  )rj  rf   r-   r  rL   rK   r)   r)   r*   test_resample_origin&  s    z*TestDataArrayResample.test_resample_originN)r   r  r  r  r  r  r  r   r  r  r  r  r  r  r  r   r  r   r5   r  r  r  r	  r  r  r)   r)   r)   r*   r    s0   
!
	:
$



(		r  c                   @  s`   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd Ze	dd Z
dd Zdd Zdd ZdS )TestDatasetResamplec           	      C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}|jddjdd}|jdddgd}t	|| t jdddd}|j
|d}|jdd}dD ]}t||}| }t|| qtjfD ]}||}t|| qd S )NrW   rA  r   r   r[   r   r   r1   r   r  r   r   rD  r[   r  rZ   Tr  r   r   r   r     )r   rR   rl  r  )rc   rd   r   r$   r%   r&   r  rl  r?   r   r  r|  r
   r   r   )	rj  rf   r(   rL   rK   Zexpected_timehowr   rs   r)   r)   r*   test_resample_and_first1  s(    


z+TestDatasetResample.test_resample_and_firstc                   s   t jdddd}tdddgtjddd	fdtjdd
dif|d t d dktj d  d<  j	ddj
dd}tj fddtd	D |d d}t|| d S )NrW   rA  r   r   r[   r   r   r1   r   r  r   r  r          @r  rZ   r   Z	min_countc                   s4   g | ],} j t|d  |d d  djdddqS )r   r   rZ   r[   r  )r?   r3   rR   )ru  ir'   r)   r*   rw  Y  s   z?TestDatasetResample.test_resample_min_count.<locals>.<listcomp>r`   )rc   rd   r   r$   r%   r&   r#   r   r   r  rR   concatr   r   )rj  rf   rL   rK   r)   r'   r*   test_resample_min_countK  s      
z+TestDatasetResample.test_resample_min_countc              	   C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}d|jd< |jddjdd}|d j}|d j}||kst	|j}|j}||kst	t
jtdd |jddd W 5 Q R X d S )NrW   rA  r   r   r[   r   r   r1   r   r  r   r  dsdatadsmetar  rZ   Tr  rD  r  r   r  )rc   rd   r   r$   r%   r&   rL  r  r   r4   r5   r   r  )rj  rf   r(   resampled_dsrL   rK   r)   r)   r*   %test_resample_by_mean_with_keep_attrsa  s(    


 z9TestDatasetResample.test_resample_by_mean_with_keep_attrsc                 C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}d|jd< |jddd j	}|j	
 d }| jtd7  _t|}t|| d S )NrW   rA  r   r   r[   r   r   r1   r   r  r   r  r  r  r  r  r  )rc   rd   r   r$   r%   r&   rL  r  r   rD  r  r  r   r   r  r   )rj  rf   r(   rL   r  rK   r)   r)   r*   test_resample_loffsetz  s    

z)TestDatasetResample.test_resample_loffsetc                 C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}d|jd< |jddjdd}|d ji ksxt	|ji kst	d S NrW   rA  r   r   r[   r   r   r1   r   r  r   r  r  r  r  rZ   Fr  rD  )
rc   rd   r   r$   r%   r&   rL  r  r   r4   rj  rf   r(   r  r)   r)   r*   &test_resample_by_mean_discarding_attrs  s    
z:TestDatasetResample.test_resample_by_mean_discarding_attrsc                 C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}d|jd< |jddjdd}|d ji ksxt	|ji kst	d S r  )
rc   rd   r   r$   r%   r&   rL  r  r  r4   r  r)   r)   r*   &test_resample_by_last_discarding_attrs  s    
z:TestDatasetResample.test_resample_by_last_discarding_attrsc                 C  s4  t d}t d}tjdddd}t t dd}t |d |d \}}t jt|td	}t||||d
d}t|j	||dd}	t|j	||dd}
t|d|id}t
||	|
|d}|dddg}|jddd}d|jkst|jdd }d|jkst|jddd}d|jks0td S )Nr2   r   rW   rA  r1   r   r  r   r  r  r@  r<   r[   rZ   r  r  r  r  r  r  r  r  r  r)   r)   r*   r    s$    

z4TestDatasetResample.test_resample_drop_nondim_coordsc              	   C  s   t jdddd}tdddgtjddd	fdtjdd
dif|d}tjtdd |	dd W 5 Q R X 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 d S )NrW   rA  r   r   r[   r   r   r1   r   r  r   r  zresample\(\) no longer supportsr   r  r   )ra   r  r`   )
rc   rd   r   r$   r%   r&   r5   r6   r   r  )rj  rf   r(   r)   r)   r*   test_resample_old_api  s    z)TestDatasetResample.test_resample_old_apic                 C  sb   t jdddd}tdtjddf|tdd}t|j	dd	
 d
 |jj	dd	
  d S )NrW   rA  i  r   r   r1   )r   r[   r   r   rZ   r   )rc   rd   r#   r   r$   r%   r&   r   r   r  r   r   )rj  r[   r(   r)   r)   r*    test_resample_ds_da_are_the_same  s     z4TestDatasetResample.test_resample_ds_da_are_the_samec                 C  st   ddd}t jdddd}tdd	d	d	gf|d
}tddddgf|d
}|jddj|dd	d}t|| d S )Nr   c                 S  s   |  d| | S r  r  rw   r)   r)   r*   rt     s    zBTestDatasetResample.test_ds_resample_apply_func_args.<locals>.funcr  r  r   r   r[   r  )r   r[   r  rZ   r  rz   )r   )rc   rd   r#   r   r  rI   r   )rj  rt   rf   r(   rK   rL   r)   r)   r*    test_ds_resample_apply_func_args  s    
z4TestDatasetResample.test_ds_resample_apply_func_argsN)r   r  r  r  r  r  r  r  r   r   r  r!  r"  r#  r)   r)   r)   r*   r  0  s   
r  c                  C  s   t jdddddddddgfidddddd	d
gddddddddgfdd} | djdd}t jdddddddddgfidddddd	d
g| jdd}t|ddg| | jdjdd}| j|jd< t	
d|jd< t|j| d S )Nr   r9   rb   r   r   r   r   r   r1   r2   r   r   group_idr   r%  r`   r   )r#   r   r=   cumsumr%  r   r=  r   r]   r$   r   rF  r)   r)   r*   test_groupby_cumsum  s$    ( r'  c                  C  s   t jdddddddddgfidddddd	d
gddddddddgfdd} | djdd}t jdddddddddgfidddddd	d
g| jdd}t|ddg| | jdjdd}| j|jd< t	
d|jd< t|j| d S )Nr   r9   rb   r   r   r   r   r   r1   r2   r   r$  r   r%  r`      )r#   r   r=   cumprodr%  r   r=  r   r]   r$   r   rF  r)   r)   r*   test_groupby_cumprod  s$    ( r*  zmethod, expected_arrayr&  r  r  g      @g      @r)  strzlist[float])r   expected_arrayr0   c              	   C  s   t jdddddddtjgfidtjdddd	id
}t|jdd| dd}t jdd|fidtjdddd	id
}t|	dg| t|j
jdd| dd}|j|jd< t|	dgj
| d S )Nr   r[   r   r   r   z
01-01-2001r   r2   r   r   Z3MrZ   r`   )r#   r   r$   r   rc   rd   r|  r  r   r=  r   r[   r]   )r   r,  r(   rL   rK   r)   r)   r*   test_resample_cumsum&  s"     
 r-  )Z
__future__r   r   r^  r$   Zpandasrc   r5   Zpandas.tseries.frequenciesr   Zxarrayr#   r   r   r   Zxarray.core.groupbyr   Zxarray.testsr   r	   r
   r   r   r   r   r   r  r+   r-   r8   rB   rM   rT   rV   rj   rl   rn   r  r  rq   ru   r|   r}   r   r   r   r   r   r%   r&   rd   r  r   Zrepr_dar   r   r   filterwarningsr   r   r   r   r
  r  r  r  r  r  r!  r%  r)  r.  r2  r>  r?  rC  rG  rI  rN  rP  rW  r[  r\  r_  r`  r  r  r'  r*  r-  r)   r)   r)   r*   <module>   s   (

	
(."

tu
	
G	
"&	   h  q E