U
    »mœdf  ã                   @   s¼   d dl mZmZ d dlZd dlmZmZ d dlm	Z
 d dlmZmZmZ d dlmZ d dlm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S )é    )ÚproductÚpermutationsN)Úassert_array_lessÚassert_allclose)Úraises)ÚinvÚeighÚnorm)Úorthogonal_procrustes)Úmatrixc                  C   s>   t j d¡ t j ddd¡} t j ddd¡}ttt| |ƒ d S )NéÒ  é   é   é   ©ÚnpÚrandomÚseedÚrandnÚassert_raisesÚ
ValueErrorr
   ©ÚAÚB© r   ú[/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/scipy/linalg/tests/test_procrustes.pyÚ)test_orthogonal_procrustes_ndim_too_large   s    r   c                  C   s6   t j d¡ t j d¡} t j d¡}ttt| |ƒ d S )Nr   r   r   r   r   r   r   Ú)test_orthogonal_procrustes_ndim_too_small   s    r   c                  C   sN   t j d¡ d} t| dƒD ].\}}t jj|Ž }t jj|Ž }ttt||ƒ qd S )Nr   ))r   r   )r   r   ©r   r   ©r   r   é   )r   r   r   r   r   r   r   r
   )ZshapesÚaÚbr   r   r   r   r   Ú)test_orthogonal_procrustes_shape_mismatch   s    r#   c            	      C   sš   t j d¡ d\} }t j | |¡}t j | |¡}t jt j t jfD ]R}| ¡ }||d< | ¡ }||d< ||f||f||ffD ]\}}ttt	||ƒ q|qBd S )Nr   )r    r   )é   r    )
r   r   r   r   ÚinfÚnanÚcopyr   r   r
   )	ÚmÚnZA_goodZB_goodZ	bad_valueZA_badZB_badr   r   r   r   r   Ú0test_orthogonal_procrustes_checkfinite_exception#   s    r*   c            
      C   s    t j d¡ d\} }tdƒD ]~}t j | |¡}t j | |¡}t||ƒ\}}t  t j d¡¡D ]<}t  t j d¡¡D ]$}t|| || ƒ\}	}t|	|ƒ qrq\qd S )Nr   r   r   )r   r   r   Úranger   r
   Úsquarer   )
r(   r)   ÚiÚA_origÚB_origZR_origÚsZA_scaleZB_scaleÚRr   r   r   Ú+test_orthogonal_procrustes_scale_invariance1   s    r2   c                  C   s¬   t j d¡ dD ]–\} }t j | |¡}t j | |¡}|| ¡ t|ƒf}|| ¡ t|ƒf}t||ƒ\}}| |¡}t||ƒD ]*\}	}
t|	|
ƒ\}}| |¡}t	||ƒ qzqd S )Nr   ©)é   r   r   )r   r4   )
r   r   r   r   Útolistr   r
   Údotr   r   )r(   r)   ZA_arrZB_arrZAsZBsZR_arrr0   ZAR_arrr   r   r1   ÚARr   r   r   Ú+test_orthogonal_procrustes_array_conversion>   s    

r8   c                  C   s
  t j d¡ dD ]ô\} }t j | |¡}t j ||¡}t|j| ƒ\}}tt|ƒ|jƒ t  ||j¡}t	||ƒ\}}tt|ƒ|jƒ t| |¡|ƒ |dt j | |¡  }	t	|	|ƒ\}
}tt|
ƒ|
jƒ |	 |¡}|	 |
¡}t
|| dd�}t
|| dd�}t||ƒ qd S )Nr   r3   g{®Gáz„?Zfro)Úord)r   r   r   r   r   ÚTr   r   r6   r
   r	   r   )r(   r)   r   ÚXÚwÚVr   r1   r0   ZA_perturbedZR_primeZnaive_approxZoptim_approxZnaive_approx_errorZoptim_approx_errorr   r   r   Útest_orthogonal_procrustesM   s$    

r>   c                 C   s   | j dd�}| | |fS )Nr   )Zaxis)Zmean)r   Úmur   r   r   Ú	_centeredn   s    r@   c            
      C   sª   t jddgddgddgddggtd�} t jddgddgddgddggtd�}t| ƒ\}}t|ƒ\}}t||ƒ\}}|t  t|ƒ¡ }|t  ||¡ | }	t|	|d	d
� d S )Néýÿÿÿr   éþÿÿÿr    ©Zdtyper$   r   r   ç:Œ0âŽyE>©Zatol©	r   ÚarrayÚfloatr@   r
   r,   r	   r6   r   )
r.   r/   r   ÚA_mur   ÚB_mur1   r0   ÚscaleÚB_approxr   r   r   Ú(test_orthogonal_procrustes_exact_examples   s    &&rM   c                  C   sP  t jddgddgddgddggtd�} t jddgddgdd	gd
dggtd�}t| ƒ\}}t|ƒ\}}t||ƒ\}}|t  t|ƒ¡ }|t  ||¡ | }	t jddgddgddgddggtd�}
t|	|
dd� d}t  t|	| ƒt|ƒ ¡}t||ƒ t||ƒ\}}|t  t|ƒ¡ }|t  ||¡ | }t  t||  ƒt|ƒ ¡}t||ƒ d S )NrA   r   rB   r    rC   é(   r$   r   iØÿÿÿr   é   iîÿÿÿiëÿÿÿé   rD   rE   gñÒÇ×ÎÜ?rF   )r.   r/   r   rI   r   rJ   r1   r0   rK   rL   ÚexpectedZexpected_disparityZAB_disparityZA_approxZBA_disparityr   r   r   Ú,test_orthogonal_procrustes_stretched_example‹   s"    &&&
rR   c                  C   sÚ   t jddgddgddgddggtd�} t jddgddgddgddggtd�}t  d	d
gd	dgd	dgddgg¡}t| ƒ\}}t|ƒ\}}t||ƒ\}}|t  t|ƒ¡ }	|	t  ||¡ | }
t|
|ƒ t|t|ƒ |ƒ d S )Nr   rB   éüÿÿÿiúÿÿÿr    rC   r$   r   gCuRãÎÁ¿g”'œ‚aå?gCuRãÎÁ?gä¯ûT6¨Ù¿gä¯ûT6¨Ù?rF   )r.   r/   ZB_standardizedr   rI   r   rJ   r1   r0   rK   rL   r   r   r   Ú(test_orthogonal_procrustes_skbio_example¡   s    &&ü
rT   )Ú	itertoolsr   r   Únumpyr   Znumpy.testingr   r   Zpytestr   r   Zscipy.linalgr   r   r	   r
   Zscipy.sparse._sputilsr   r   r   r#   r*   r2   r8   r>   r@   rM   rR   rT   r   r   r   r   Ú<module>   s"   	!