a
    ߙfbK                     @   s   d dl Z d dlZd dlZd dlmZ d dlmZ d dl	m
Z
 G dd dZG dd dZG d	d
 d
ZG dd dZG dd dZdS )    N)assert_array_equal)units)NUMPY_LT_1_21_1c                   @   s8   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d ZdS )TestQuantityArrayCopyz?
    Test whether arrays are properly copied/used in place
    c                 C   sb   t d}tj|ddd}tj|ddd}d|d< |d j|d ksHJ |d j|d ks^J d S )Ng     @@km/sF)copyT      r   )nparangeuQuantityvalue)selfvZq_nocopyZq_copy r   Nlib/python3.9/site-packages/astropy/units/tests/test_quantity_array_methods.pytest_copy_on_creation   s    
z+TestQuantityArrayCopy.test_copy_on_creationc                 C   s   t tddd}|t jt j }t|j|jks<J |t j	t j }t|j|jksdJ dt j	 t j |d< |d j|d jksJ d S )N      ?      Y@r   r   r   )
r   r   r	   r
   tomsallr   km)r   qq2q3r   r   r   test_to_copies   s    z$TestQuantityArrayCopy.test_to_copiesc                 C   s^   t tdd}|j}t|j|jks.J dt j t j |d< |d j|d jksZJ d S )Nr   m/sr   r   )	r   r   r	   r
   Zsir   r   r   r   )r   r   r   r   r   r   test_si_copies#   s
    z$TestQuantityArrayCopy.test_si_copiesc                 C   s   t tdd}|dd }dt j t j |d< |d |d ksFJ |ddtjf }dt j t j |d	< t| |ksJ dS )
zCheck that [keys] work, and that, like ndarray, it returns
        a view, so that changing one changes the other.

        Also test that one can add axes (closes #1422)
        r   r   
      r   r   Ni)r    r   )	r   r   r	   r
   r   r   Znewaxisr   flatten)r   r   Zq_selr   r   r   r   test_getitem_is_view*   s    z*TestQuantityArrayCopy.test_getitem_is_viewc                 C   sf  t tdddd}|j}|d dt j t j ks<J t|dd tdt j t j kshJ d	d
 |jD }tt |t dd
 |j	jD |j
ksJ dt j t j |d< |d dt j t j ksJ |d dt j t j ksJ dt j t j |d< |d dt j t j ks&J |d dt j t j ksDJ |d dt j t j ksbJ d S )N      "@   r             @r             @c                 S   s   g | ]}|qS r   r   ).0Z_qr   r   r   
<listcomp>A       z3TestQuantityArrayCopy.test_flat.<locals>.<listcomp>c                 S   s   g | ]}|qS r   r   )r*   Z_ar   r   r   r+   C   r,   r   )r(   r(   g       )r   r   r	   r
   reshapeflatr   r   r   r   unitr   )r   r   Zq_flatZq_flat_listr   r   r   	test_flat9   s    ,
zTestQuantityArrayCopy.test_flatN)	__name__
__module____qualname____doc__r   r   r   r#   r1   r   r   r   r   r      s   	r   c                   @   sb   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	e
jjejdkoNedddd ZdS )TestQuantityReshapeFuncszTest different ndarray methods that alter the array shape

    tests: reshape, squeeze, ravel, flatten, transpose, swapaxes
    c                 C   s^   t dtj }|dd}t|tjs,J |j|jks<J t |j	|j	ddksZJ d S N      @r%   r(   )
r	   r
   r   r   r.   
isinstancer   r0   r   r   )r   r   Z	q_reshaper   r   r   test_reshapeU   s
    z%TestQuantityReshapeFuncs.test_reshapec                 C   s^   t dddtj }| }t|tjs0J |j|jks@J t 	|j
|j
 ksZJ d S )Nr8         )r	   r
   r.   r   r   squeezer9   r   r0   r   r   )r   r   Z	q_squeezer   r   r   test_squeeze\   s
    z%TestQuantityReshapeFuncs.test_squeezec                 C   s^   t dddtj }| }t|tjs0J |j|jks@J t 	|j
|j
 ksZJ d S r7   )r	   r
   r.   r   r   Zravelr9   r   r0   r   r   )r   r   Zq_ravelr   r   r   
test_ravelc   s
    z#TestQuantityReshapeFuncs.test_ravelc                 C   s^   t dddtj }| }t|tjs0J |j|jks@J t 	|j
|j
 ksZJ d S r7   )r	   r
   r.   r   r   r"   r9   r   r0   r   r   )r   r   Z	q_flattenr   r   r   test_flattenj   s
    z%TestQuantityReshapeFuncs.test_flattenc                 C   s^   t dddtj }| }t|tjs0J |j|jks@J t 	|j
|j
 ksZJ d S r7   )r	   r
   r.   r   r   Z	transposer9   r   r0   r   r   )r   r   Zq_transposer   r   r   test_transposeq   s
    z'TestQuantityReshapeFuncs.test_transposec                 C   sh   t ddddtj }|dd}t|tjs6J |j|jksFJ t 	|j
|j
ddksdJ d S )Nr8   r%   r<   r(   r   )r	   r
   r.   r   r   Zswapaxesr9   r   r0   r   r   )r   r   Z
q_swapaxesr   r   r   test_swapaxesx   s
    z&TestQuantityReshapeFuncs.test_swapaxesZbigzNumpy GitHub Issue 19153)reasonc                    s   t ddddtj }|j t dks0J  j|u s>J  jdksLJ  j	dksZJ  fdd	 D d
 }|d dksJ |d dksJ  
 }t|| ksJ t|tjsJ t ||rJ dS )z9While ``flat`` doesn't make a copy, it changes the shape.r8   r%   r<   r(   r;   )r   r   r   r   c                    s   g | ]} j  jfqS r   )indexcoords)r*   xZqfr   r   r+      r,   zATestQuantityReshapeFuncs.test_flat_attributes.<locals>.<listcomp>r-      )r(   r   r<   N)r	   r
   r.   r   r   r/   lenbaserE   rD   r   r   r"   r9   r   Zmay_share_memory)r   r   Z
endindicesZq_flat_copyr   rG   r   test_flat_attributes   s    z-TestQuantityReshapeFuncs.test_flat_attributesN)r2   r3   r4   r5   r:   r>   r?   r@   rA   rB   pytestZmarkZxfailsys	byteorderr   rK   r   r   r   r   r6   O   s   r6   c                   @   s0  e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"d# Zd$d% Zd&d' Zd(d) Zd*d+ Zd,d- Zd.d/ Zd0d1 Zd2d3 Zd4d5 Zd6d7 Zd8d9 Zd:d; Z d<d= Z!d>d? Z"d@dA Z#dBdC Z$dDdE Z%dFdG Z&dHdI Z'dJS )KTestQuantityStatsFuncsz$
    Test statistical functions
    c                 C   sJ   t g dtj }tt |dtj  tt j|dddgtj  d S )Nr   r)         @      @r8   @T)keepdims)r	   arrayr   r   r   meanr   q1r   r   r   	test_mean   s    z TestQuantityStatsFuncs.test_meanc                 C   sN   t g dtj }dtj }t j||d}||u s8J |dtj ksJJ d S )NrP         ?outrS   )r	   rU   r   r   r   rV   )r   rX   qiqi2r   r   r   test_mean_inplace   s
    
z(TestQuantityStatsFuncs.test_mean_inplacec                 C   sJ   t ddgtj }tt |dtj  t|jddddgtj  d S )Nr   r)         ?TZaxisrT   )r	   rU   r   r   r   stdrW   r   r   r   test_std   s    zTestQuantityStatsFuncs.test_stdc                 C   sB   t ddgtj }dtj }t j||d |dtj ks>J d S )Nr   r)   rZ   r[   r`   )r	   rU   r   r   r   rc   r   rX   r]   r   r   r   test_std_inplace   s    
z'TestQuantityStatsFuncs.test_std_inplacec                 C   sR   t ddgtj }tt |dtjd   t|jddddgtjd   d S )Nr   r)         ?r(   r   Trb   )r	   rU   r   r   r   varrW   r   r   r   test_var   s    zTestQuantityStatsFuncs.test_varc                 C   sF   t ddgtj }dtj }t j||d |dtjd  ksBJ d S )Nr   r)   rZ   r[   rg   r(   )r	   rU   r   r   r   rh   re   r   r   r   test_var_inplace   s    
z'TestQuantityStatsFuncs.test_var_inplacec                 C   s0   t g dtj }t |dtj ks,J d S )NrP   rQ   )r	   rU   r   r   medianrW   r   r   r   test_median   s    z"TestQuantityStatsFuncs.test_medianc                 C   sB   t g dtj }dtj }t j||d |dtj ks>J d S )NrP   rZ   r[      )r	   rU   r   r   r   rk   re   r   r   r   test_median_inplace   s    
z*TestQuantityStatsFuncs.test_median_inplacec                 C   s0   t g dtj }t |dtj ks,J d S )NrP   r   )r	   rU   r   r   minrW   r   r   r   test_min   s    zTestQuantityStatsFuncs.test_minc                 C   sB   t g dtj }dtj }t j||d |dtj ks>J d S )NrP   rZ   r[   r   )r	   rU   r   r   r   ro   re   r   r   r   test_min_inplace   s    
z'TestQuantityStatsFuncs.test_min_inplacec                 C   s*   t g dtj }t |dks&J d S )N)r8   r)   rQ   rR   r8   r<   )r	   rU   r   r   ZargminrW   r   r   r   test_argmin   s    z"TestQuantityStatsFuncs.test_argminc                 C   s0   t g dtj }t |dtj ks,J d S )NrP   r8   )r	   rU   r   r   maxrW   r   r   r   test_max   s    zTestQuantityStatsFuncs.test_maxc                 C   sB   t g dtj }dtj }t j||d |dtj ks>J d S )NrP   rZ   r[   r8   )r	   rU   r   r   r   rs   re   r   r   r   test_max_inplace   s    
z'TestQuantityStatsFuncs.test_max_inplacec                 C   s*   t g dtj }t |dks&J d S )N)rR   r)   rQ   rR   r8   rm   )r	   rU   r   r   ZargmaxrW   r   r   r   test_argmax   s    z"TestQuantityStatsFuncs.test_argmaxc                 C   s^   t g dtj tj }|ddtj tj }t |t g dtj tj ksZJ d S )NrP           @rZ   r)   rQ   rR   rx   )r	   rU   r   r   r   clipMmr   r   rX   Zc1r   r   r   	test_clip   s    z TestQuantityStatsFuncs.test_clipc                 C   s   t g dtj tj }|jddtj tj |d}t |t g dtj tj ks^J dtj tj |d< t |j	|j	ksJ d S )NrP   rw   rx   r[   ry   r    r   )
r	   rU   r   r   r   rz   r{   r   Zmmr   r|   r   r   r   test_clip_inplace   s
    (z(TestQuantityStatsFuncs.test_clip_inplacec                 C   s4   t g dtj tj }t | |ks0J d S )NrP   )r	   rU   r   r   r   r   ZconjrW   r   r   r   	test_conj   s    z TestQuantityStatsFuncs.test_conjc                 C   s0   t g dtj }t |dtj ks,J d S )NrP   rR   )r	   rU   r   r   ptprW   r   r   r   test_ptp   s    zTestQuantityStatsFuncs.test_ptpc                 C   sB   t g dtj }dtj }t j||d |dtj ks>J d S )NrP   rZ   r[   rR   )r	   rU   r   r   r   r   re   r   r   r   test_ptp_inplace   s    
z'TestQuantityStatsFuncs.test_ptp_inplacec                 C   s   t g dtj }t t |t g dtj ks<J t t j|ddt j|jddtj ksjJ t |jdd|jjddtj ksJ d S )Ng?5^I?g/$@g/$
@r<   r(   r%   r(   decimals)r	   rU   r   kgr   roundr   rW   r   r   r   
test_round  s    (
z!TestQuantityStatsFuncs.test_roundc                 C   s\   t g dtj }t dtj }|jd|d}||u s>J t |jdd|ksXJ d S )Nr   r%   r(   )r   r\   r   )r	   rU   r   r   zerosr   r   r   )r   rX   r]   ar   r   r   test_round_inplace	  s
    z)TestQuantityStatsFuncs.test_round_inplacec                 C   s   t g dtj }t | dtj ks0J t t |dtj ksNJ t g dg dgtj }t |dt g dtj ksJ t t |dt g dtj ksJ d S )Nr   r)   r8   r$   )rQ   rR   r$   )r   r   r   r   )rR   r8         $@)r	   rU   r   r   r   sumr   r   rX   r   r   r   r   test_sum  s    (zTestQuantityStatsFuncs.test_sumc                 C   sB   t g dtj }dtj }t j||d |dtj ks>J d S )Nr   rZ   r[   r$   )r	   rU   r   r   r   r   re   r   r   r   test_sum_inplace  s    
z'TestQuantityStatsFuncs.test_sum_inplacec                 C   s   t g dtj }t | t g dtj ks:J t t |t g dtj ksbJ t g dtj }t | t g dtj ksJ t t |t g dtj ksJ d S )Nr<   r(   r;   r<   r%   	   )rm   rH   r   )rm   r      )r	   rU   r   r   r   cumsumr   r   r   r   r   test_cumsum   s    &(&z"TestQuantityStatsFuncs.test_cumsumc                 C   sz   t g dtj }t dtj }t j||d t |t g dtj ksTJ |}|j|d t ||ksvJ d S )Nr   r%   r[   r   )r	   rU   r   r   Zonesr   r   r   )r   rX   r]   r   r   r   r   test_cumsum_inplace*  s    "z*TestQuantityStatsFuncs.test_cumsum_inplacec                 C   s   t ddt jgtj }t | dtj ks4J t t |dtj ksRJ t t jddgdt jdggtj }t |dt g dtj ksJ t t |dt g dtj ksJ d S )Nr   r)         @rR   r$   r   )r   rR   r   )r	   rU   nanr   r   r   nansumr   r   r   r   r   test_nansum3  s    $(z"TestQuantityStatsFuncs.test_nansumc                 C   s   t ddt jgtj }dtj }|j|d}||u s:J |t |j|j ksTJ dtj }t j||d}||u sxJ |t |j|j ksJ d S )Nr   r)   rZ   r[   )	r	   rU   r   r   r   r   r   r   r0   )r   rX   r]   Zqoutr^   Zqout2r   r   r   test_nansum_inplace=  s    

z*TestQuantityStatsFuncs.test_nansum_inplacec                 C   s   t g dtj }ttj}|  W d    n1 s>0    Y  ttj}t | W d    n1 st0    Y  t g dtd }| dtd ksJ t |dtd ksJ d S )Nr   )r   rQ   rR   r<   g      N@)	r	   rU   r   r   rL   raises
UnitsErrorZprodUnitr   rX   excr   r   r   r   	test_prodJ  s    &(z TestQuantityStatsFuncs.test_prodc                 C   s   t g dtj }ttj}|  W d    n1 s>0    Y  ttj}t | W d    n1 st0    Y  t g dtd }t 	| t g dtd ksJ t 	t |t g dtd ksJ d S )Nr   )r%   rm   rH   r<   )r%      <   )
r	   rU   r   r   rL   r   r   Zcumprodr   r   r   r   r   r   test_cumprodV  s    &(*z#TestQuantityStatsFuncs.test_cumprodc                 C   sf   t g dtj }t | t g dtj ks:J t t |t g dtj ksbJ d S Nr   r)   rQ   r   r   )r	   rU   r   r   r   ZdiffrW   r   r   r   	test_diffb  s    &z TestQuantityStatsFuncs.test_diffc                 C   sf   t g dtj }t | t g dtj ks:J t t |t g dtj ksbJ d S r   )r	   rU   r   r   r   Zediff1drW   r   r   r   test_ediff1dh  s    &z#TestQuantityStatsFuncs.test_ediff1dc                 C   sf   t g dtj }t g dtj }||}|jt |j|jksLJ |jtjtj ksbJ d S )Nr   )r   rQ   rR   r8   )r	   rU   r   r   r   dotr   r0   )r   rX   r   r   r   r   r   test_dot_methn  s
    
z$TestQuantityStatsFuncs.test_dot_methc                 C   s8   t ddgddggtj }t |dtj ks4J d S )Nr   r)   r   rQ   rR   )r	   rU   r   r   tracer   r   r   r   r   test_trace_funcv  s    z&TestQuantityStatsFuncs.test_trace_funcc                 C   s~   t ddgddggtj }| dtj ks2J tdtj}t ddgddggtj }|j|d |dtj kszJ d S )	Nr   r)   r   rQ   rR   r8   r[   r$   )r	   rU   r   r   r   r   r   )r   rX   contr   r   r   r   test_trace_meth{  s    z&TestQuantityStatsFuncs.test_trace_methc                 C   sL   t dtj }t t |dtj dtj t g dtj ksHJ d S )Nr    r%   r;   
r   r   r   r   rQ   rR   r8   r8   r8   r8   )r	   r
   r   r   r   rz   rU   r   r   r   r   test_clip_func  s    z%TestQuantityStatsFuncs.test_clip_funcc                 C   s   t g dtj }t dtj }|dtj dtj }t |dtj dtj |ksbJ t dtj }|jdtj dtj |d t ||ksJ d S )Nr   r    r%   r;   r[   )	r	   rU   r   r   r
   rz   r   r   r   )r   ZexpectedrX   r   r   r   r   r   test_clip_meth  s    &z%TestQuantityStatsFuncs.test_clip_methN)(r2   r3   r4   r5   rY   r_   rd   rf   ri   rj   rl   rn   rp   rq   rr   rt   ru   rv   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   rO      sJ   

	
rO   c                   @   s`   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Zdd ZdS )TestArrayConversionz'
    Test array conversion methods
    c                 C   s  t jtg dt jt j td}|dd|j ks:J |	dd |ddt j t j ksdJ |	ddt j
 t j  |ddt j t j ksJ tt( |	ddt j t j  W d    n1 s0    Y  tt |	  W d    n1 s0    Y  d|d< |d dt j t j ks6J dt j
 t j |d< |d dt j t j kshJ tt$ dt j t j |d< W d    n1 s0    Y  d S )Nr   dtyper<   r(   i  d   rZ   )r   r   r	   rU   r   r   intitemr0   ZitemsetcmrL   r   	TypeError
ValueErrorrW   r   r   r   	test_item  s     "6(zTestArrayConversion.test_itemc                 C   s   t g dtj tj }|ddtj tj ks8J t|dt ddgtj tj ksdJ |dd t |dt dd	g|j ksJ |d
dtj	 tj  |
d
dtj tj ksJ d S )Nr   r<   r(   )r   r(   r%   )r<   r(   )r%   rm   i  i  r   i  rH   )r	   rU   r   r   r   Ztaker   Zputr0   r   r   rW   r   r   r   test_take_put  s    ,(z!TestArrayConversion.test_take_putc                 C   s  t g dg dgtj tj }| }d|d< |j|jksBJ |d jdksTJ dtj tj |d< t	|
 dd	 jt g d
ksJ d|dddf< t	|
 dd	 jt g dksJ ttj dtj |d< W d   n1 s0    Y  tt dd}dtj tj |d< |d jdks<J d|d< |d dksVJ t j|d< t |d stJ t j|d< t |d sJ dS )zrTest that setitem changes the unit if needed (or ignores it for
        values where that is allowed; viz., #2695))r   r)   r   )rQ   rR   r8   g     @)r   r   r   r$   r   Nr%   )r$   r$   r$   g     @@ra   )r'   r'   r$   r    )r<   r<   r   r   rH   r           )r	   rU   r   r   r   r   r0   r   r{   r   r"   rL   r   r   r   r   r
   r   infZisinfr   Zisnan)r   r   rX   r   r   r   r   
test_slice  s(    "((,

zTestArrayConversion.test_slicec                 C   sF   t g dtj tj }|d t |dtj tj ksBJ d S )Nr   r(   i  )r	   rU   r   r   r   Zfillr   rW   r   r   r   	test_fill  s    
zTestArrayConversion.test_fillc                 C   s  t g dtj tj }|d}|j|jks4J t|j|jdksNJ |	  |j|jksfJ |
t g d}|j|jksJ t|j|j
t g dksJ t ddgddggtj tj }| }|j|jksJ t|j|j ksJ d S )Nr   r(   )TTFFr<   r%   rm   )r	   rU   r   r   r   repeatr0   r   r   sortcompressZdiagonalr   r   r   r   test_repeat_compress_diagonal  s    
$"z1TestArrayConversion.test_repeat_compress_diagonalc                 C   s   t jg dt jdtj tj }|t j}t|dr:J |tj	}|j
d u sTJ t j|}|j
|j
u spJ |tj	}|j|ju sJ d S )Nr   r   r0   )r	   rU   int64r   r   r   viewZndarrayhasattrr   Z_unitZmaZMaskedArrayr0   )r   rX   r   r   Zq4Zq5r   r   r   	test_view  s     zTestArrayConversion.test_viewc                 C   s   t jjdd}|ddddf jtj tj \}}}t |dddf tj tj | }t|tj	slJ |dtj tj ksJ dS )zT
        Regression test for https://github.com/astropy/astropy/issues/2003
        )r    r&   )sizeNr<   rm   r   )
r	   ZrandomZuniformTr   r   r   r   r9   r   )r   r   rF   yztotalr   r   r   test_slice_to_quantity  s
    (&z*TestArrayConversion.test_slice_to_quantityc                 C   s   t jg dt jdtj tj }| }|j|jks8J t|j	|j	 ksPJ |
t j}t||kslJ |jt jks|J |jt jdd}| jt jdd}|j|jksJ t| |ksJ d S )Nr   r   r   )offsetrm   )r	   rU   r   r   r   r   Zbyteswapr0   r   r   ZastypeZfloat64r   ZgetfieldZint32)r   rX   r   Zq2aZq2br   r   r   !test_byte_type_view_field_changes  s     z5TestArrayConversion.test_byte_type_view_field_changesc                 C   s|   t g dtj tj }| }t|dr0J |  |ddg}t|drTJ t	||
tjjddgksxJ d S )N)r   rR   r)   rQ   r0   rw   i	  )r	   rU   r   r   r   Zargsortr   r   Zsearchsortedr   r   Zdimensionless_unscaledr   )r   rX   ir   r   r   	test_sort  s    zTestArrayConversion.test_sortc                 C   s  t g dtj tj }tt |g d W d    n1 sH0    Y  tt |	  W d    n1 sz0    Y  tt |
  W d    n1 s0    Y  tt |  W d    n1 s0    Y  tt |d W d    n1 s0    Y  tt |d W d    n1 sJ0    Y  tt |  W d    n1 s~0    Y  d S )Nr   )r   r   r<   r   za.a)r	   rU   r   r   r   rL   r   NotImplementedErrorZchoosetolistZtostringtobytesZtofiledumpdumpsrW   r   r   r   test_not_implemented  s    ,&&&**z(TestArrayConversion.test_not_implementedN)r2   r3   r4   r5   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r     s   	
r   c                   @   s(   e Zd ZdZdd Zdd Zdd ZdS )	TestRecArrayzcRecord arrays are not specifically supported, but we should not
    prevent their use unnecessarilyc                 C   s*   t t dddjdd | _d S )Ng      (@rm   r%   zf8,f8,f8r   )r	   rU   r
   r.   r   r=   ra)r   r   r   r   setup1  s    zTestRecArray.setupc                 C   s:   t | jt j}t|d d j| jd d ks6J d S )Nr(   )r   r   r   r   r	   r   r   r   Zqrar   r   r   test_creation5  s    zTestRecArray.test_creationc                 C   s4   t | jt j}|d |d< |d |d ks0J d S )Nr(   r<   )r   r   r   r   r   r   r   r   test_equality9  s    zTestRecArray.test_equalityN)r2   r3   r4   r5   r   r   r   r   r   r   r   r   -  s   r   )rM   rL   Znumpyr	   Znumpy.testingr   Zastropyr   r   Zastropy.utils.compatr   r   r6   rO   r   r   r   r   r   r   <module>   s   BK   