U
    iâËd  ã                   @   sü   d dl Zd dlmZmZmZ d dlmZ d dlm	Z	 d dl
mZ d dlZejZejded� dd	¡Ze ¡ Zejd
d�Zejded� dd¡ddd…ddd…f Zdd„ Zdd„ Zdd„ Zdd„ ZegZG dd„ deƒZG dd„ deƒZedkrøe ¡  dS )é    N)Úfloat32ÚjitÚnjit)Ú	Vectorize)ÚTypingError)ÚTestCaseéP   )Údtypeé   é
   ÚF)Úorderi@  é   é   é   c                 C   s   | | S ©N© )ÚaÚbr   r   úX/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/numba/tests/npyufunc/test_ufunc.pyÚadd   s    r   c                 C   s   | | | | S r   r   )r   r   ÚcÚdr   r   r   Úadd_multiple_args   s    r   c                 C   s0   d}t | jd ƒD ]}|| | ||  7 }q|S )Ng        r   )ÚrangeÚshape)r   r   ÚresultÚir   r   r   Ú
gufunc_add   s    r   c                 C   s   t |jƒD ]}|  |¡}q
|S r   )r   ÚndimÚreduce)ÚufuncÚargr   r   r   r   Úufunc_reduce!   s    r#   c                   @   sD   e 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S )Ú
TestUFuncsc                 G   sª   |t f|žŽ }|  tttƒ¡ | ¡ }||jf}| j|||ƒ|| |d� | jt||ƒt |¡|d� | j| |¡tj  |¡|d� | j| 	||¡tj  	||¡|d� dS )zTest ufunc attributes©ÚmsgN)
r   r   Úbuild_ufuncr   ÚassertPreciseEqualr#   ÚnpÚsumÚ
accumulateÚouter)ÚselfÚclsr   r   ÚargsÚ
vectorizerr!   Úinfor   r   r   Ú_test_ufunc_attributes2   s    
ÿz!TestUFuncs._test_ufunc_attributesc           	      C   sX   |t ƒ}| tttttƒ¡ | ¡ }||jf}| j|||||ƒ|| | | |d� dS )zTest multiple argsr%   N)r   r   r   r'   r   r(   )	r-   r.   r   r   r   r   r0   r!   r1   r   r   r   Ú_test_broadcasting?   s
    
zTestUFuncs._test_broadcastingc              
   C   s|   t D ]}|  |td td ¡ qt D ]}|  |tt¡ q$t D ]:}|  |td d …tjd d …f ttjd d …d d …f ¡ q<d S ©Nr   )Úvectorizersr2   r   r   r)   Únewaxis©r-   Úvr   r   r   Útest_ufunc_attributesH   s    ÿz TestUFuncs.test_ufunc_attributesc                 C   s¸   t D ]&}|  |td td td td ¡ qt D ]}|  |tttt¡ q0t D ]f}|  |td d …tjd d …f ttjd d …d d …f td d …tjd d …f ttjd d …d d …f ¡ qLd S r4   )r5   r3   r   r   r   r   r)   r6   r7   r   r   r   Útest_broadcastingQ   s    $2 ÿzTestUFuncs.test_broadcastingc                 C   sd   t D ]Z}|tƒ}| tttƒ¡ | ¡ }ttjd d …tjtjd d …f }|  |t|ƒt| ¡ qd S r   )	r5   r   r   r'   r   r)   r6   r(   r   )r-   r8   r0   r!   Zbroadcasting_br   r   r   Útest_implicit_broadcastingZ   s     ÿz%TestUFuncs.test_implicit_broadcastingc           	   
   C   s–   t  d¡}d|j_g }d}| tdd�t|f¡ | tdd�tdf¡ |D ]F\}}}dd	„ }|  |¡�}||ƒ|ƒ W 5 Q R X |  	|t
|jƒ¡ qJd S )
Nr   Fz<ufunc 'sin' called with an explicit output that is read-onlyT)Znopython)Zforceobjzoutput array is read-onlyc                 S   s    t  | j| j¡}t  || ¡ d S r   )r)   Úonesr   r	   Úsin)Úxr   r   r   r   Útesto   s    zBTestUFuncs.test_ufunc_exception_on_write_to_readonly.<locals>.test)r)   r<   ÚflagsZ	writeableÚappendr   r   Ú
ValueErrorÚassertRaisesÚassertInÚstrÚ	exception)	r-   ÚzÚtestsÚexpectÚdecÚexcr&   r?   Úraisesr   r   r   Ú)test_ufunc_exception_on_write_to_readonlyd   s    
ÿz4TestUFuncs.test_ufunc_exception_on_write_to_readonlyc              	   C   st   t dd„ ƒ}|  |t d¡dƒt dd¡ tj¡¡ |  t¡�}|t d¡dƒ W 5 Q R X d}|  |t	|j
ƒ¡ d S )Nc                 S   s$   |dkrd }n
t  d¡}t  | |¡S )Nr   é   )r)   r<   r   )r>   ÚyrG   r   r   r   Úinner{   s    
z5TestUFuncs.test_optional_type_handling.<locals>.innerrN   é   é   é   z(expected array(float64, 1d, C), got None)r   r(   r)   ÚarangeZastypeZfloat64rC   Ú	TypeErrorrD   rE   rF   )r-   rP   rL   r&   r   r   r   Útest_optional_type_handlingx   s    
ÿz&TestUFuncs.test_optional_type_handlingN)
Ú__name__Ú
__module__Ú__qualname__r2   r3   r9   r:   r;   rM   rV   r   r   r   r   r$   0   s   			
r$   c                   @   s   e Zd Zdd„ Zdd„ ZdS )ÚTestUFuncsMiscc                 C   sH   t dd„ ƒ}tjtjfD ]*}|dƒ}| |¡}||ƒ}|  ||¡ qd S )Nc                 S   s
   t  | ¡S r   )r)   Zexp2©r>   r   r   r   Úfoo•   s    z%TestUFuncsMisc.test_exp2.<locals>.foor   ©r   r)   Zint8Zuint16Zpy_funcr(   ©r-   r\   Útyr>   ÚexpectedÚgotr   r   r   Ú	test_exp2“   s    

zTestUFuncsMisc.test_exp2c                 C   sH   t dd„ ƒ}tjtjfD ]*}|dƒ}| |¡}||ƒ}|  ||¡ qd S )Nc                 S   s
   t  | ¡S r   )r)   Úlog2r[   r   r   r   r\   ¡   s    z%TestUFuncsMisc.test_log2.<locals>.foor   r]   r^   r   r   r   Ú	test_log2Ÿ   s    

zTestUFuncsMisc.test_log2N)rW   rX   rY   rb   rd   r   r   r   r   rZ   �   s   rZ   Ú__main__)Únumpyr)   Znumbar   r   r   Znumba.np.ufuncr   Znumba.core.errorsr   Znumba.tests.supportr   Zunittestr	   rT   Zreshaper   Úcopyr   r   r   r   r   r   r#   r5   r$   rZ   rW   Úmainr   r   r   r   Ú<module>   s(   *ÿ	`