U
    Hvf(                     @   sH  d Z ddlZddlZddlZddlmZmZmZmZ ddl	m
Z
mZmZmZmZmZ ddlZddlZddlmZmZmZmZmZmZmZmZmZmZ ddlmZ ddlZddlm Z  dd	 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&G dd dZ'dd Z(dd Z)dd Z*G dd dZ+G dd dZ,dS ) z, Test functions for linalg.matfuncs module

    N)arrayidentitydotsqrt)assert_array_almost_equalassert_allcloseassert_assert_array_lessassert_array_equalassert_warns)
funmsignmlogmsqrtmfractional_matrix_powerexpmexpm_frechet	expm_condnorm
khatri_rao)_matfuncs_inv_ssq)minimizec                  C   s:   t jddddgddddgddddgddddggtd} | S )aW  
    Return the test matrix from Experiment (1) of [1]_.

    References
    ----------
    .. [1] Awad H. Al-Mohy and Nicholas J. Higham (2012)
           "Improved Inverse Scaling and Squaring Algorithms
           for the Matrix Logarithm."
           SIAM Journal on Scientific Computing, 34 (4). C152-C169.
           ISSN 1095-7197

    g3d?g     L@r   gRal!A?gQI?g^?dtype)npr   float)A r   D/tmp/pip-unpacked-wheel-96ln3f52/scipy/linalg/tests/test_matfuncs.py%_get_al_mohy_higham_2012_experiment_1   s    



r   c                   @   s,   e Zd Zdd Zdd Zdd Zdd Zd	S )
	TestSignMc              
   C   s   t dddddgdddd	d
gddddd
gddddd
gdddddgg}t dddddgdddddgdd d!d"d gd#d$d%d&d$gd'd(d)d*d(gg}t|}t|| d S )+N333333=@3333338     `Q@fffffH@      @ffffff"@      20             $      @      4333333#333333#@     9.皙#@333333      2@3333332@       @g'@g'Rqgy.@ge@5@g@ȽgUЛ'?g9\g+%gRQgQ?gGzgffffffgר#gw?î?gk#egC,Ng@g42
gF(@gcs@@)r   r   r   )selfaZcrrr   r   r   	test_nils0   s    zTestSignM.test_nilsc                 C   s@   t ddddgddddgddddgddddgg}t|dd d S )N           r   Fdispr   r   r7   r8   r   r   r   test_defective1>   s    0zTestSignM.test_defective1c              
   C   sT   t dddddgdddd	d
gddddd
gddddd
gdddddgf}t|dd d S )Nr!   r"   r#   r$   r%   r&   r'   r(   r)   r*   r+   r,   r-   r.   r/   r0   r1   r2   r3   r4   r5   r6   Fr=   r?   r@   r   r   r   test_defective2C   s    zTestSignM.test_defective2c                 C   s   t dddddddgdddddddgdddddddgdddddddgdddddddgdddddddgdddddddgg}t|d	d
 d S )Nr*         9@r;               $@      @r6         .@Fr=   r?   r@   r   r   r   test_defective3M   s    zTestSignM.test_defective3N)__name__
__module____qualname__r:   rA   rB   rH   r   r   r   r   r    .   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d Z
dd Zdd Zdd ZdS )TestLogMc                 C   s   t dddddddgdddddddgdddddddgdddddddgdddddddgdddddddgdddddddgg}td	d
 d | }t|dd d S )Nr*   rC   r;   rD   rE   rF   r6   rG      g@y                Fr=   )r   r   r   )r7   r8   mr   r   r   r:   [   s    zTestLogM.test_nilsc                 C   s2   t  }t|dd\}}t|}t||ddd d S )NFr=   g-C6
?+=rtolatol)r   r   r   r   )r7   r   A_logminfoA_round_tripr   r   r   *test_al_mohy_higham_2012_experiment_1_logmg   s    z3TestLogM.test_al_mohy_higham_2012_experiment_1_logmc                 C   s>   t  }t|tjdd\}}t|}ttj||ddd  d S )NFr=   h㈵>rO   rP   )r   r   r   logr   r   Zallclose)r7   r   Z
A_funm_logrT   rU   r   r   r   .test_al_mohy_higham_2012_experiment_1_funm_logo   s    z7TestLogM.test_al_mohy_higham_2012_experiment_1_funm_logc                 C   s   t jd tddD ]}t j||}t dddD ]n}|| }t j|}d||}t	|dd	\}}|
|}	t|	| t|dd	\}
}t|
}t|||d
 q6qd S )N  r<         	   zM:{0} eivals:{1}Fr=   )err_msg)r   randomseedrangerandnlogspacelinalgeigvalsformatr   r   r   r   r   )r7   n
M_unscaledscaleMWr_   M_sqrtmrT   M_sqrtm_round_tripM_logmZM_logm_round_tripr   r   r   test_round_trip_random_floatw   s    

z%TestLogM.test_round_trip_random_floatc                 C   s~   t jd tddD ]b}t j||dt j||  }t dddD ].}|| }t|dd	\}}t|}t|| qHqd S 
NrZ   r<   r[                 ?r\   r]   r^   Fr=   )	r   r`   ra   rb   rc   rd   r   r   r   )r7   rh   ri   rj   rk   ro   rT   M_round_tripr   r   r   test_round_trip_random_complex   s     z'TestLogM.test_round_trip_random_complexc                 C   s   d}ddgddggddgddggddgddggddgddggfD ]}t j|}ttdd |D   tj|td}t|d	d
\}}t|j	j
|k tj|td}t|d	d
\}}t|j	j
|k tj|td }t|d	d
\}}t|j	j
|k q@d S )NFDGr<   r         c                 s   s   | ]}|j p|jd k V  qdS r   Nimagreal.0wr   r   r   	<genexpr>   s     zFTestLogM.test_logm_type_preservation_and_conversion.<locals>.<genexpr>r   Fr=   )scipyre   rf   r   anyr   r   r   r   r   charcomplex)r7   complex_dtype_charsmatrix_as_listrl   r   rS   rT   r   r   r   *test_logm_type_preservation_and_conversion   s"    z3TestLogM.test_logm_type_preservation_and_conversionc                 C   s   dddgdddgdddgg}t tfD ]h}tj||d}tj|}tdt|j	
 k  t|dd\}}tt|jtj tt|| q$d S )Nr<   ry   r   {Gz?Fr=   )r   r   r   r   r   re   rf   r   absoluter}   sumr   
issubdtyper   Zinexactr   r   )r7   rk   dtXr   YrT   r   r   r   test_complex_spectrum_real_logm   s    z(TestLogM.test_complex_spectrum_real_logmc                 C   sj   ddgddggddgddggfD ]D}t tfD ]6}tj||d}t|dd\}}tt|jtj q,q d S )Nr<   r   r   Fr=   )	r   r   r   r   r   r   r   r   Zcomplexfloating)r7   rk   r   r   rS   rT   r   r   r   test_real_mixed_sign_spectrum   s    z&TestLogM.test_real_mixed_sign_spectrumc                 C   sv   t ddgddgg}t ddgddgg}||j||jfD ]4}tj}t|t|dd\}}t|}t	||dd q<d S )Nr   rr   r<   Fr=   rO   rR   )
r   r   ZasarrayTr   ZLogmExactlySingularWarningr   r   r   r   )r7   r   Brk   expected_warningLrT   Er   r   r   test_exactly_singular   s    zTestLogM.test_exactly_singularc                 C   sB   t dgg}tj}t|t|dd\}}t|}t||dd d S )Ng0.++Fr=   rO   r   )r   r   r   ZLogmNearlySingularWarningr   r   r   r   )r7   rk   r   r   rT   r   r   r   r   test_nearly_singular   s
    zTestLogM.test_nearly_singularc                 C   s  ddgddgg}dt jd gt j d dgg}tt||dd tt||dd ddgdd	gg}dt j d d
t j gdd	t j d gg}tt||dd tt||dd ddgdd	gg}dt j d dgdd	t j d gg}tt||dd tt||dd d S )Nr   r<   r         ?rO   r   rr   r]   y             ry   )r   pir   r   r   )r7   r   r   r   r   r   &test_opposite_sign_complex_eigenvalues   s    *$z/TestLogM.test_opposite_sign_complex_eigenvaluesN)rI   rJ   rK   r:   rV   rY   rp   rt   r   r   r   r   r   r   r   r   r   r   rL   Y   s   
	rL   c                   @   s   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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)S )*	TestSqrtMc                 C   sn   t jd tddD ]R}t j||}t dddD ]0}|| }t|dd\}}||}t|| q6qd S )	NrZ   r<   r[   r\   r]   r^   Fr=   	r   r`   ra   rb   rc   rd   r   r   r   r7   rh   ri   rj   rk   rm   rT   rn   r   r   r   rp      s    
z&TestSqrtM.test_round_trip_random_floatc                 C   s   t jd tddD ]d}t j||dt j||  }t dddD ]0}|| }t|dd	\}}||}t|| qHqd S rq   r   r   r   r   r   rt      s     
z(TestSqrtM.test_round_trip_random_complexc                 C   s   d}t |}tddddgd|ddgdd|dgddddgg}tddddgd|ddgdd|dgddddgg}|jd }tt||| t|d|dd }tt||| t|dddd }tt||| d S )	Ng      ?      ?r   r<   r   Fr>   	blocksizery   )r   r   shaper   r   r   )r7   eser8   sarh   Zesar   r   r   test_bad   s$    






zTestSqrtM.test_badc                 C   s  d}ddgddggddgddggddgddggddgddggddgddggfD ]}t j|}ttdd |D   tj|td}t|d	d
\}}t|j	j
|k tj|td}t|d	d
\}}t|j	j
|k tj|td }t|d	d
\}}t|j	j
|k qNd S )Nru   r<   r   ry   rz   c                 s   s   | ]}|j p|jd k V  qdS r{   r|   r   r   r   r   r   #  s     zHTestSqrtM.test_sqrtm_type_preservation_and_conversion.<locals>.<genexpr>r   Fr=   )r   re   rf   r   r   r   r   r   r   r   r   r   r7   r   r   rl   r   A_sqrtmrT   r   r   r   +test_sqrtm_type_preservation_and_conversion  s$    z5TestSqrtM.test_sqrtm_type_preservation_and_conversionc                 C   s   d}ddgddggddgddggdddgdddgdddggfD ]}t j|}ttdd |D  tj|td}t|dd	\}}t|j	j
|k tj|td}t|dd	\}}t|j	j
|k q>d S )
Nru   r<   r   r   c                 s   s   | ]}|j p|jd k V  qdS r{   r|   r   r   r   r   r   =  s     zVTestSqrtM.test_sqrtm_type_conversion_mixed_sign_or_complex_spectrum.<locals>.<genexpr>r   Fr=   )r   re   rf   r   r   r   r   r   r   r   r   r   r   r   r   r   9test_sqrtm_type_conversion_mixed_sign_or_complex_spectrum4  s    zCTestSqrtM.test_sqrtm_type_conversion_mixed_sign_or_complex_spectrumc                 C   s   t jd tddD ]v}t j||dt j||  }t|d|d\}}t|t j	|d tddD ] }t|d|d\}}t|| qjqd S )	NrZ   r<      rr   Fr   ry   
   )
r   r`   ra   rb   randrc   r   r   re   matrix_power)r7   rh   r   ZA_sqrtm_defaultrT   r   ZA_sqrtm_newr   r   r   test_blocksizesI  s     zTestSqrtM.test_blocksizesc                 C   sH   t  }t|dd\}}||}t||dd tt|t| d S )NFr=   rW   rQ   )r   r   r   r   r   tril)r7   r   r   rT   rU   r   r   r   %test_al_mohy_higham_2012_experiment_1T  s
    
z/TestSqrtM.test_al_mohy_higham_2012_experiment_1c              	   C   sj   t tfD ]\}tjddddgddddgddddgddddgg|d}t|dd\}}tt|  qd S )Nr   rz   r   Fr=   )intr   r   r   r   r   isnanall)r7   r   r   r   rT   r   r   r   test_strict_upper_triangular\  s    



z&TestSqrtM.test_strict_upper_triangularc                 C   s   t tfD ]}tjdddgdddgdddgg|d}tjdddgdddgdddgg|d}t||| t|dd\}}tt|	  qd S )Nr   r<   r   Fr=   )
r   r   r   r   r
   r   r   r   r   r   )r7   r   r   r   ZB_sqrtmrT   r   r   r   test_weird_matrixg  s$    zTestSqrtM.test_weird_matrixc                 C   s:   t jd t jdd}t|dd}t||| d S )NrZ   rz   Tr=   )r   r`   ra   r   r   r   r   )r7   r   r   r   r   r   	test_dispx  s    zTestSqrtM.test_dispc                 C   sL   ddgddgg}ddgddgg}t t|||dd	 t t||dd	 d S )
N               @r]   r                       ?      ?ry         ?      rO   r   )r   r   r   r   r7   rk   Rr   r   r   r     s    z0TestSqrtM.test_opposite_sign_complex_eigenvaluesc              
   C   s   t ddddgddddgddddgddddgg}t tdddtdgddddgddddgtdddtdgg}tt |||dd tt||dd d S )Nr<   r   r   rO   r   )r   r   r   r   r   r   r   r   r   r   test_gh4866  s    




zTestSqrtM.test_gh4866c                 C   sP   t dddg}t tdddg}tt |||dd tt||dd d S )Nry   r<   r   rO   r   )r   Zdiagr   r   r   r   r   r   r   r   test_gh5336  s    zTestSqrtM.test_gh5336c                 C   s@   t d}t d}tt |||dd tt||dd d S )N)ry   ry   rO   r   )r   zerosr   r   r   r   r   r   r   test_gh7839  s    

zTestSqrtM.test_gh7839c                 C   s   t jdt jd}t|jt jks$tt jdt jd}t|jt jksHtt jdt jd}t|jt j	ksltt jdt j
d}t|jt jkstd S Nr   r   r   )r   r   Zuint8r   r   float16AssertionErrorZuint16Zuint32float32Zuint64float64r7   rk   r   r   r   -test_data_size_preservation_uint_in_float_out  s    z7TestSqrtM.test_data_size_preservation_uint_in_float_outc                 C   s   t jdt jd}t|jt jks$tt jdt jd}t|jt jksHtt jdt jd}t|jt j	ksltt jdt j
d}t|jt jkstd S r   )r   r   int8r   r   r   r   int16int32r   int64r   r   r   r   r   ,test_data_size_preservation_int_in_float_out  s    z6TestSqrtM.test_data_size_preservation_int_in_float_outc                 C   s   t jddgddggt jd}t|jt jks0tt jddgddggt jd}t|jt jks`tt jddgddggt jd}t|jt jkstt jddgddggt j	d}t|jt j
kstd S )Nry   r]   r   r   )r   r   r   r   r   	complex64r   r   r   r   
complex128r   r   r   r   +test_data_size_preservation_int_in_comp_out  s    z5TestSqrtM.test_data_size_preservation_int_in_comp_outc                 C   s   t jdt jd}t|jt jks$tt jdt jd}t|jt jksHtt jdt jd}t|jt jkslttt drt jdt j	d}t|jt j	kstd S )Nr   r   float128)
r   r   r   r   r   r   r   r   hasattrr   r   r   r   r   .test_data_size_preservation_float_in_float_out  s    
z8TestSqrtM.test_data_size_preservation_float_in_float_outc                 C   s   t jddgddggt jd}t|jt jks0tt jddgddggt jd}t|jt jks`tt jddgddggt jd}t|jt j	kstt
t drt
t drt jddgddggt jd}t|jt jkstd S )Nry   r]   r   r   r   r   
complex256)r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   -test_data_size_preservation_float_in_comp_out  s    z7TestSqrtM.test_data_size_preservation_float_in_comp_outc                 C   s   t jddgddggt jd}t|jt jks0ttt drt jddgddggt jd}t|jt jksjtt jddgddggt jd}t|jt jkstd S )Nr   r]   r   r   r   r   )	r   r   r   r   r   r   r   r   r   r   r   r   r   ,test_data_size_preservation_comp_in_comp_out  s    
z6TestSqrtM.test_data_size_preservation_comp_in_comp_outN)rI   rJ   rK   rp   rt   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r      s(   

r   c                   @   sj   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d Z
ejjdddd Zdd ZdS )TestFractionalMatrixPowerc                 C   s   t jd tddD ]v}tddD ]f}t j||dt j||  }t dddD ]2}|| }t|d| }t j||}t	|| qVq$qd S )NrZ   r<      rr   r\   r]   r^   
r   r`   ra   rb   rc   rd   r   re   r   r   r7   prh   ri   rj   rk   ZM_rootrs   r   r   r   rt     s     z8TestFractionalMatrixPower.test_round_trip_random_complexc                 C   s   t jd tddD ]d}tddD ]T}t j||}t dddD ]2}|| }t|d| }t j||}t	|| qDq$qd S )NrZ   r<   r   r\   r]   r^   r   r   r   r   r   rp     s    z6TestFractionalMatrixPower.test_round_trip_random_floatc                 C   s   t jd dD ]}tdD ]}t j||dt j||  }t|d}t j|d}t|| t|d}t j|d}t|| t|d	}t j|d
}t|| qqd S )NrZ   )ry   rz   r   r   rr   皙?r   gigffffff@   )	r   r`   ra   rb   rc   r   re   r   r   )r7   rh   irk   ZM_one_fifthrs   r   r   r   r   r   (test_larger_abs_fractional_matrix_powers  s     




zBTestFractionalMatrixPower.test_larger_abs_fractional_matrix_powersc                 C   s   t jd d}t|D ]}tdd}t j }t tdd}t j||}tdrr|dt j||  }|| }t||}t	|dd	\}}	t
|| }
t||
 qd S )
NrZ      r<   r   r\   )TFrr   Fr=   )r   r`   ra   rb   	randrangerc   expchoicer   r   r   r   )r7   nsamplesr   rh   r   Zmatrix_scaler   A_powerrS   rT   ZA_power_expm_logmr   r   r   test_random_matrices_and_powers  s    


z9TestFractionalMatrixPower.test_random_matrices_and_powersc           	      C   s   t  }t|tjdd\}}t|dd\}}t|d}t|d}t|| t	|| t	|| dD ]D}t||}t|d| }t	||dd t	t
|dt
|d qbd S )NFr=   r   )r   g?r<   r   r   )r   r   r   r   r   r   Z_remainder_matrix_powerr   r
   r   r   )	r7   r   ZA_funm_sqrtrT   r   ZA_rem_powerr   r   rU   r   r   r   r   &  s    




z?TestFractionalMatrixPower.test_al_mohy_higham_2012_experiment_1c                 C   sj   t jd t jddt jd  D ]<}tdD ].}t||}|t |  d }t|| q4q(d S )NrZ   r   rr   r   r<   )	r   r`   ra   rc   rb   r   Z_briggs_helper_functionZexp2r   )r7   r8   kZ
x_observedZ
x_expectedr   r   r   test_briggs_helper_function:  s     z5TestFractionalMatrixPower.test_briggs_helper_functionc                 C   s   d}ddgddggddgddggddgddggddgddggfD ]}t j|}ttdd |D   dD ]~}tj|td	}t||}t|j	j
|k tj|td	}t||}t|j	j
|k tj|td	 }t||}t|j	j
|k qlq@d S )
Nru   r<   r   ry   rz   c                 s   s   | ]}|j p|jd k V  qdS r{   r|   r   r   r   r   r   P  s     zRTestFractionalMatrixPower.test_type_preservation_and_conversion.<locals>.<genexpr>333333r   gffffff
@r   )r   re   rf   r   r   r   r   r   r   r   r   r   r7   r   r   rl   r   r   r   r   r   r   %test_type_preservation_and_conversionB  s$    


z?TestFractionalMatrixPower.test_type_preservation_and_conversionc                 C   s   d}ddgddggddgddggdddgdddgdddggfD ]}t j|}ttdd |D  dD ]T}tj|td}t||}t|j	j
|k tj|td}t||}t|j	j
|k qhq>d S )	Nru   r<   r   r   c                 s   s   | ]}|j p|jd k V  qdS r{   r|   r   r   r   r   r   n  s     z`TestFractionalMatrixPower.test_type_conversion_mixed_sign_or_complex_spectrum.<locals>.<genexpr>r   r   )r   re   rf   r   r   r   r   r   r   r   r   r   r   r   r   r   3test_type_conversion_mixed_sign_or_complex_spectrume  s    

zMTestFractionalMatrixPower.test_type_conversion_mixed_sign_or_complex_spectrumzToo unstable across LAPACKs.reasonc                 C   s   ddgddggddgddggddgddggdddgdddgdddggfD ]r}t tfD ]d}tj||d}dD ] }t||}tt|  qjd	D ]&}t||}t|d| }t|| qqTqHd S )
Nr   r<   ry   rz   r[   r   r   )gffffffr   r   g)r   gzG?)	r   r   r   r   r   r   r   r   r   )r7   r   Znewtyper   r   r   rU   r   r   r   test_singular~  s    

z'TestFractionalMatrixPower.test_singularc                 C   sN   ddgddgg}ddgddgg}t t|||dd	 t t|d
|dd	 d S )Nr   r]   r   r   r   ry   r   rO   r   r   )r   r   r   r   r   r   r   r   r     s    z@TestFractionalMatrixPower.test_opposite_sign_complex_eigenvaluesN)rI   rJ   rK   rt   rp   r   r   r   r   r   r   pytestmarkZxfailr   r   r   r   r   r   r     s   #
r   c                   @   s4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )TestExpMc                 C   s2   t ddgddgg}tt|ddgddgg d S )Nr;   r   r<   )r   r   r   r@   r   r   r   	test_zero  s    zTestExpM.test_zeroc                 C   s"   t d}t|ttjgg d S )Nr<   )r   r   r   r   r   )r7   eltr   r   r   test_single_elt  s    zTestExpM.test_single_eltc                 C   s$   t d}t|}|jdks td S )Nr   r   r   )r   r   r   sizer   )r7   r   resultr   r   r   test_empty_matrix_input  s    
z TestExpM.test_empty_matrix_inputc                 C   s   t j}tddgddgg}|d d d|  }|d d | }tt|t||g|d |gg t|t jjjdks|t	t|t j
jjdkst	d S )Nr<   r]   ry   rv   f)r   r   r   r   r   Zastyper   r   r   r   r   )r7   r   r8   ZaaZbbr   r   r   test_2x2_input  s    "zTestExpM.test_2x2_inputc                 C   sx  t j}t jddgddggddgddggddgddggddgddggddgddgggdd	}t |d d d
|  |d d | g|d d d | |d d d
|  ggdd|d
   d|d
  d  d|d
  d dd|d
    g|d
 d dd|d
    dd|d
   |d
 d  ggdd|  |d d  dd|  d|d  d  gdd
|  |d d
  dd|  d|d  d  ggdd|d
   d|d  d  dd|d
   d|d  d  gdd|d
   d|d  d  dd|d
   d|d  d  ggdd
|  d| d
  dd
|  d| d
  gdd
|  d| d
  dd
|  d| d
  ggg}tt|| d S )Nr<   r]   rz   r   r   r\   rv   )orderry   rM   r   r[   )r   r   r   r   r   )r7   r   r8   Za_resr   r   r   test_nx2x2_input  s6    "&>622>>..zTestExpM.test_nx2x2_inputN)rI   rJ   rK   r   r   r  r  r  r   r   r   r   r     s
   	r   c                   @   sJ   e Zd Zdd Zdd Zdd Zdd Zejj	ejj
d	d
dd ZdS )TestExpmFrechetc           	      C   s   t jddddgddddgd	d	ddgd	d	ddggtd
}t jddgddggtd
}t jddgddggtd
}tj|}tj|d ddd f }i ddiddifD ]*}t||f|\}}t|| t|| qd S )Nr<   ry   rz   r]   r   r[   rM   r   r   r   methodSPSblockEnlarge)r   r   r   r   re   r   r   r   )	r7   rk   r   r   expected_expmexpected_frechetkwargsobserved_expmobserved_frechetr   r   r   test_expm_frechet  s4    




z!TestExpmFrechet.test_expm_frechetc                 C   sJ  t jddddgddddgd	d	ddgd	d	ddggtd
}t jddgddggtd
}t jddgddggtd
}tj|d}ddddddddg}t|d d |dd  }|D ]\}}tjjj| }	tjjj| }
d|	|
  }|| }|| }|| }|| }tj	|}tj	|d ddd f }t
||\}}t|| t|| qd S )Nr<   ry   rz   r]   r   r[   rM   r   r   r   r^            r   r   )r   r   r   r   re   r   zipZ_expm_frechetZell_table_61r   r   r   )r7   Z
M_original
A_original
E_originalA_original_norm_1Zselected_m_listZm_neighbor_pairsmambZell_aZell_btarget_norm_1rj   rk   r   r   r  r  r  r  r   r   r   test_small_norm_expm_frechet  sH    




z,TestExpmFrechet.test_small_norm_expm_frechetc              	   C   s  t jjt jjt jjt jjf}d}t|D ]}t|}td}t	dd}|||fd}|||fd}t
j|d}	||	 }
|
| }|
| }t t ||gt t ||gg}t
j|}t
j|d ||d f }t||\}}t||dd t||d	d q(d S )
Nd   r   ry      r   r<   gHj>r   gHz>)r   r`   uniformnormalZstandard_cauchyexponentialrb   r   expovariater   r   re   r   vstackZhstackZ
zeros_liker   r   r   )r7   ZrfuncsZntestsr   Zrfuncr  rh   r  r  r  rj   r   r   rk   r  r  r  r  r   r   r   	test_fuzz  s2    

zTestExpmFrechet.test_fuzzc                 C   s~   t jddgddggtd}t jddgdd	ggtd}tj|d
 t||dd\}}t||dd\}}t|| t|| d S )Ng]؉??gvP?g^?guÝ?r   gS,?g| @gJt?g~3?r<   r
  r	  r  )r   r   r   r   re   r   r   r   )r7   r   r   sps_expmsps_frechetblockEnlarge_expmblockEnlarge_frechetr   r   r   test_problematic_matrix  s2      
  

z'TestExpmFrechet.test_problematic_matrixzthis test is deliberately slowr   c                 C   sd   d}t jj||fd}t jj||fd}t||dd\}}t||dd\}}t|| t|| d S )Ni  r  r
  r&  r  )r   r`   r"  r   r   )r7   rh   r   r   r'  r(  r)  r*  r   r   r   test_medium_matrix0  s      
  

z"TestExpmFrechet.test_medium_matrixN)rI   rJ   rK   r  r  r%  r+  r   r   slowskipr,  r   r   r   r   r    s   !r  c           
      C   sL   t || j}t|}|| ||  }t| | }t|| ||  }	|	 S N)r   reshaper   r   r   )
r   A_normr   X_normepsr   Zp_normperturbationX_primeZscaled_relative_errorr   r   r   _help_expm_cond_search?  s    r6  c                 C   s   | t j|t j|   S r/  )r   re   r   )r   r   r   r   r   _normalized_likeH  s    r7  c                 C   s(   | |}| || }t || t | S r/  )r   )r  r   r4  r   r5  r   r   r   _relative_errorL  s    r8  c                   @   s4   e Zd Zdd Zdd Zdd Zejjdd Z	d	S )
TestExpmConditionNumberc                 C   s@   t jd tddD ]$}t j||}t|}td| qd S )NrZ   r<   r]   r   )r   r`   ra   rb   rc   r   r	   )r7   rh   r   kappar   r   r   test_expm_cond_smokeS  s
    z,TestExpmConditionNumber.test_expm_cond_smokec              	   C   sH   t ddddgddddgddd	d
gddddgg}t|}td| d S )Ng(gٙBgx@g  =g ZdBr   gzNV7g@g  2pgᙥT g?W[E@gT!g yhG)r   r   r   r	   )r7   r   r:  r   r   r   test_expm_bad_condition_numberZ  s    



z6TestExpmConditionNumber.test_expm_bad_condition_numberc                 C   s   t jd t jddddD ]$}t |gg}tt|t| qt jddddD ]$}t |gg}tt|t| qRt	dD ]*}t j
d	d	}tt|t |d
  qd S )N90  r  r   r  )numr   ry   r   r<   r   )r   r`   ra   Zlinspacer   r   r   absrd   rb   rc   r   )r7   xr   r   r   r   r   test_univariated  s    z'TestExpmConditionNumber.test_univariatec              	   C   sF  t jd d}d}t|D ]"}t jdd}t j||}tj|}t	|}tj|}t
|}	tt|||||}
t || }t|
|dd}|j}|
|}|tt ||j| }tt	||}t|| |  tdD ]D}|tt jj|j | }tt|t| tt	||}t|| qt|dd|  | |	  qd S )	Nr=  rW   r   ry   r   zL-BFGS-Br&  r<   )r   r`   ra   rb   randintrc   r   re   r   r   r   	functoolspartialr6  Zonesr   r@  r7  r0  r   r8  r   r	   )r7   r3  r   r   rh   r   r1  r   r2  r:  r  guessoutZxoptZyoptZp_bestZp_best_relerrjZp_randZp_rand_relerrr   r   r   test_expm_cond_fuzzp  s<        z+TestExpmConditionNumber.test_expm_cond_fuzzN)
rI   rJ   rK   r;  r<  rA  r   r   r-  rH  r   r   r   r   r9  R  s
   
r9  c                   @   s,   e Zd Zdd Zdd Zdd Zdd Zd	S )
TestKhatriRaoc                 C   s~   t tddgddggtddgddgg}t|tdd	gdd
gddgddgg t tddgtddg}t|jd d S )Nr<   ry   rz   r]   r   r[   rM   r      r  r            )r]   ry   )r   r   r
   r   emptyr   r7   r8   br   r   r   
test_basic  s    zTestKhatriRao.test_basicc              	   C   sP   t t< tdddgdddgg}tddgddgg}t|| W 5 Q R X d S )Nr<   ry   rz   r]   r   r[   r   Zraises
ValueErrorr   r   rO  r   r   r   test_number_of_columns_equality  s    
z-TestKhatriRao.test_number_of_columns_equalityc              	   C   s   t t, tdddg}tdddg}t|| W 5 Q R X t t6 tdddg}tdddgdddgg}t|| W 5 Q R X t t6 tdddgddd	gg}tdddg}t|| W 5 Q R X d S )
Nr<   ry   rz   r]   r   r[   rM   r   r^   rR  rO  r   r   r   test_to_assure_2d_array  s$    z%TestKhatriRao.test_to_assure_2d_arrayc                    sf   t ddgddgg t ddgddggt }t fd	d
tjd D j}t|| d S )Nr<   ry   rz   r]   r   r[   rM   r   c                    s0   g | ](}t  d d |f d d |f qS r/  )r   Zkron)r   r   r8   rP  r   r   
<listcomp>  s   z@TestKhatriRao.test_equality_of_two_equations.<locals>.<listcomp>)r   r   r   r$  rb   r   r   r
   )r7   Zres1Zres2r   rV  r   test_equality_of_two_equations  s    

z,TestKhatriRao.test_equality_of_two_equationsN)rI   rJ   rK   rQ  rT  rU  rX  r   r   r   r   rI    s   rI  )-__doc__r`   rC  Znumpyr   r   r   r   r   Znumpy.testingr   r   r   r	   r
   r   r   Zscipy.linalgr   r   r   r   r   r   r   r   r   r   r   r   Zscipy.linalg._expm_frechetZscipy.optimizer   r   r    rL   r   r   r   r  r6  r7  r8  r9  rI  r   r   r   r   <module>   s4    0+  v :0v	E