U
    »mœdÉ(  ã                   @   s(  d dl Z d dlZd dlZd dlZd dlZd dlmZmZ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Z d dlZd dlmZ dd„ Zd	d
„ Zdd„ Zdd„ Zejjej d¡dd„ ƒƒZejj e j! "d¡oæe #ej$¡j%dk dd�G dd„ dƒƒZ&ejj'dd�dd„ ƒZ(dd„ Z)dd„ Z*dS )é    N)Úassert_equalÚassert_Úassert_allclose)Ú_sparsetoolsÚ
coo_matrixÚ
csr_matrixÚ
csc_matrixÚ
bsr_matrixÚ
dia_matrix)Úsupported_dtypes)Úcheck_free_memory)Úraisesc                 C   s   | d d d S )z6
    Wrap an integer to the interval [-128, 127].
    é€   é   © )Únr   r   ú\/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/scipy/sparse/tests/test_sparsetools.pyÚint_to_int8   s    r   c                   C   s   t ttjƒ d S ©N)Úassert_raisesÚMemoryErrorr   Ztest_throw_errorr   r   r   r   Útest_exception   s    r   c                     sœ   d} d‰d‰t t ˆˆg¡ƒ‰g ‰G ‡‡‡‡fdd„dtjƒ‰ ‡ fdd„t| ƒD ƒ}|D ]}| ¡  qX|D ]}| ¡  qjˆD ]}tt 	| 
¡ dk¡ƒ q|d S )	Né
   éd   é   c                       s   e Zd Z‡ ‡‡‡fdd„ZdS )ztest_threads.<locals>.Workerc                    sT   ˆ   ¡ }tˆƒD ]4}t ˆˆˆ jˆ jˆ jˆ jˆ jˆ j|j|j|j¡ qˆ |¡ d S r   )ÚcopyÚranger   Zcsr_plus_csrÚindptrÚindicesÚdataÚappend)ÚselfÚbÚj©ÚaÚbresr   Úniterr   r   Úrun(   s          ýz test_threads.<locals>.Worker.runN)Ú__name__Ú
__module__Ú__qualname__r(   r   r$   r   r   ÚWorker'   s   r,   c                    s   g | ]
}ˆ ƒ ‘qS r   r   )Ú.0Ú_)r,   r   r   Ú
<listcomp>1   s     z test_threads.<locals>.<listcomp>é   )r   ÚnpÚonesÚ	threadingÚThreadr   ÚstartÚjoinr   ÚallÚtoarray)ZnthreadsÚthreadsÚthreadr"   r   )r,   r%   r&   r   r'   r   Útest_threads   s    


r;   c                  C   sZ   t D ]P} t ddgddgg¡ | ¡}t|| d�}t| d¡ ¡ |d d …d d…f ƒ qd S )Né   r0   é   é   ©Údtyper   )r   r1   ÚarrayÚastyper   r   Zgetcolr8   )r@   Úadr%   r   r   r   Ú!test_regression_std_vector_dtypes;   s    rD   z!Can't create large array for testc                  C   sŽ   t  t j¡jd } td|  d d ƒ t j| t jd�}t j| t jd�}t j| t jd�}d|d< t|||ffdd	d
�}| ¡ }t	|dggƒ d S )Nr<   é	   ç    €„.Ag      à?r?   r>   éÿÿÿÿ)r<   r<   F)Úshaper   )
r1   ÚiinfoÚint32Úmaxr   ÚzerosÚint8r   r8   r   )ZnnzÚrowÚcolr   ÚsÚdr   r   r   Útest_nnz_overflowF   s    rR   Úlinuxé   ztest requires 64-bit Linux)Úreasonc                	   @   sò   e Zd ZdZdZdd„ Zdd„ Zdd„ Zej	j
d	d
„ ƒZej	j
dd„ ƒZejdej	jd�ejdej	jd�dddejdej	jd�gZej	j
ej	 de¡dd„ ƒƒZej	j
ej	 de¡dd„ ƒƒZdd„ Zdd„ Zdd„ Zdd „ Zd!d"„ Zd#d$„ Zd%S )&ÚTestInt32Overflowa  
    Some of the sparsetools routines use dense 2D matrices whose
    total size is not bounded by the nnz of the sparse matrix. These
    routines used to suffer from int32 wraparounds; here, we try to
    check that the wraparounds don't occur any more.
    iPÃ  c                 C   s^   | j d t tj¡jkst‚zttj 	dd¡ƒ}W n t
k
rL   tj}Y nX td| ƒ d S )Nr0   ZPYTEST_XDIST_WORKER_COUNTÚ1i¸  )r   r1   rI   rJ   rK   ÚAssertionErrorÚintÚosÚenvironÚgetÚ
ValueErrorÚinfr   )r!   Zparallel_countr   r   r   Úsetup_methodg   s    zTestInt32Overflow.setup_methodc                 C   s   t  ¡  d S r   )ÚgcÚcollect©r!   r   r   r   Úteardown_methods   s    z!TestInt32Overflow.teardown_methodc                 C   s€   | j }t d|d g¡}t d|d g¡}tjddgtjd�}t|||ffƒ}| ¡ }t|d dƒ t|d dƒ ~t ¡  d S ©Nr   r<   r0   r?   ©r   r   )rG   rG   )	r   r1   rA   rM   r   Ztodenser   r`   ra   )r!   r   Úir#   r   ÚmÚrr   r   r   Útest_coo_todensev   s    z"TestInt32Overflow.test_coo_todensec           
      C   s¸   | j }t d|d g¡}t d|d g¡}tjddgtjd�}t|||ffƒ}tj||ftjd�}tttfD ]<}||ƒ}| 	|¡}	t
|	d dƒ t
|	d dƒ ~	t ¡  ql~t ¡  d S rd   )r   r1   rA   rM   r   r2   r   r   r	   Údotr   r`   ra   )
r!   r   rf   r#   r   rg   r"   ZsptypeÚm2rh   r   r   r   Útest_matvecsˆ   s    

zTestInt32Overflow.test_matvecsc                 C   s€   | j }tj||ftjd�}t |¡}t||f||fd�}tj|jd tjd�}| |¡}t|d t	|ƒƒ ~~~~~t
 ¡  d S )Nr?   )rH   r<   r   )r   r1   r2   rM   Úaranger
   rH   rj   r   r   r`   ra   )r!   r   r   Úoffsetsrg   Úvrh   r   r   r   Útest_dia_matvecž   s    


z!TestInt32Overflow.test_dia_matvecZmatmat)ZmarksZmatvecsZmatvecÚdiagonalÚsort_indicesÚ	transposeÚopc                    s:   ‡ fdd„}t  ¡  ztˆ d| ƒ|ƒ W 5 t  ¡  X d S )Nc                     sd   ˆ j } tjd| | ftjd�}tjddgtjd�}tjdgtjd�}t|||f| | fdd�}~~~|S ©Nr<   r?   r   F)Ú	blocksizer   )r   r1   r2   rM   rA   rJ   r	   ©r   r   r   r   rg   rb   r   r   Ú
get_matrix¹   s    z6TestInt32Overflow.test_bsr_1_block.<locals>.get_matrixÚ_check_bsr_©r`   ra   Úgetattr©r!   rt   rx   r   rb   r   Útest_bsr_1_block²   s
    	z"TestInt32Overflow.test_bsr_1_blockc                    s:   ‡ fdd„}t  ¡  ztˆ d| ƒ|ƒ W 5 t  ¡  X d S )Nc                     sb   ˆ j } tj| | dftjd�}tjd| gtjd�}tj| tjd�}t|||f| dfdd�}~~~|S ru   )r   r1   r2   rM   rA   rJ   rm   r	   rw   rb   r   r   rx   Ð   s    z6TestInt32Overflow.test_bsr_n_block.<locals>.get_matrixry   rz   r|   r   rb   r   Útest_bsr_n_blockÈ   s
    	z"TestInt32Overflow.test_bsr_n_blockc                 C   s<   |ƒ }| j }| tj|dftjd�¡}t|d t|ƒƒ d S )Nr0   r?   re   ©r   rj   r1   r2   rM   r   r   ©r!   rg   r   rh   r   r   r   Ú_check_bsr_matvecsß   s    z$TestInt32Overflow._check_bsr_matvecsc                 C   s:   |ƒ }| j }| tj|ftjd�¡}t|d t|ƒƒ d S )Nr?   r   r   r€   r   r   r   Ú_check_bsr_matvecç   s    z#TestInt32Overflow._check_bsr_matvecc                 C   s(   |ƒ }| j }| ¡ }t|t |¡ƒ d S r   )r   rq   r   r1   r2   r€   r   r   r   Ú_check_bsr_diagonalï   s    z%TestInt32Overflow._check_bsr_diagonalc                 C   s   |ƒ }|  ¡  d S r   )rr   ©r!   rg   r   r   r   Ú_check_bsr_sort_indices÷   s    z)TestInt32Overflow._check_bsr_sort_indicesc                 C   s   |ƒ }|  ¡  d S r   )rs   r„   r   r   r   Ú_check_bsr_transposeü   s    z&TestInt32Overflow._check_bsr_transposec                 C   sr   |ƒ }| j }ttj|dftjd�|jd dfd�}| |¡ ~ttjd|ftjd�d|jd fd�}| |¡ d S )Nr0   r?   r<   )rv   r   )r   r	   r1   r2   rM   rv   rj   )r!   rg   r   rk   r   r   r   Ú_check_bsr_matmat  s    &
&z#TestInt32Overflow._check_bsr_matmatN)r)   r*   r+   Ú__doc__r   r_   rc   ri   ÚpytestÚmarkÚslowrl   rp   ÚparamZxslowZ_bsr_opsZparametrizer}   r~   r�   r‚   rƒ   r…   r†   r‡   r   r   r   r   rV   [   s8   

ûrV   z/64-bit indices in sparse matrices not availablec                  C   sŽ   d} | d t  t j¡jkst‚t| d d d ƒ t j| ft jd�}t j| d t jd�}t j	| t jd�}t
|||fƒ}|j}tt|j|ƒ d S )Nl   4s	j r0   é   r=   rF   r?   r<   )r1   rI   Zint64rK   rX   r   r2   rM   rm   rL   r   ÚTr   ÚRuntimeErrorrj   )r   r   r   r   r%   r"   r   r   r   Útest_csr_matmat_int64_overflow  s    r�   c                  C   sÐ  t tjtjd gddggtd�} tjddgtd�}tD �]’}tD �]†}d||f }t |tj¡rp|  ¡  	|¡}n| j
 ¡  	|¡}t |tj¡rž| ¡  	|¡}n(tjdd	�� |j
 ¡  	|¡}W 5 Q R X |tjkrÜ|tjk�stjd
tjd�}tttjdd|j|j|j||ƒ	 t |tj¡�r,t |tj¡�rLt |tj¡�s|t |tj¡�r|tjd
tjd�}tttjdd|j|j|j||ƒ	 tjd
t ||¡d�}t dd|j|j|j||¡ t|t | ¡ |¡|d� q@q6d S )Ny              ð?r=   r>   r?   y      p@      ð?l        z(%r, %r)Úignore)Úinvalid)r0   r0   )Úerr_msg)r   r1   ÚpiÚcomplexrA   r   Z
issubdtypeZcomplexfloatingr   rB   ÚrealZerrstateZbool_rL   r   r]   r   Z
csr_matvecr   r   r   Zfloat64Zresult_typer   rj   r8   )Za0Zb0Za_dtypeZb_dtypeÚmsgr%   r"   Úcr   r   r   Útest_upcast!  sV     

      ÿÿþý      ÿr™   c                  C   sz   t  d¡} dddg}t|  d¡|fdƒ}t|  d¡|fdƒ}t  d¡}t| |¡dd	d
dgƒ t| |¡dd	d
dgƒ d S )N)r=   r>   rG   r   r<   z<f8)r>   r>   z>f8r>   r=   é   é   )r1   r2   r
   rB   rm   r   rj   )rQ   rn   r%   r"   ro   r   r   r   Útest_endiannessH  s    


rœ   )+ÚsysrZ   r`   r3   Únumpyr1   Znumpy.testingr   r   r   Zscipy.sparser   r   r   r   r	   r
   Zscipy.sparse._sputilsr   Zscipy._lib._testutilsr   r‰   r   r   r   r   r;   rD   rŠ   r‹   Zxfail_on_32bitrR   ZskipifÚplatformÚ
startswithr@   ZintpÚitemsizerV   Úskipr�   r™   rœ   r   r   r   r   Ú<module>   s4    
$ÿ 3
'