o
    Eb>                    @   s  d Z ddlZddlmZmZmZmZmZmZ ddl	Z	ddl	m
Z ddlm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mZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z. dd	l/m0Z0 dd
l1m2Z2 ddl3m4Z4 ddl5m6Z6 G dd dZ7G dd dZ8G dd dZ9d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@G d!d" d"ZAG d#d$ d$ZBG d%d& d&ZCd'd( ZDd)d* ZEdS )+z8
Test functions for multivariate normal distributions.

    N)assert_allcloseassert_almost_equalassert_array_almost_equalassert_equalassert_array_lessassert_)raises   )check_distribution_rvs)_PSD_lnB_cho_inv_batchmultivariate_normal_frozen)multivariate_normalmultivariate_hypergeommatrix_normalspecial_ortho_grouportho_grouprandom_correlationunitary_group	dirichletbetawishartmultinomial
invwishartchi2invgammanormuniformks_2sampkstestbinom	hypergeommultivariate_tcauchy
normaltest)romb)multigammaln)check_random_state_property)patchc                   @   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)d* Zd+d, Zd-d. Zd/d0 Zd1S )2TestMultivariateNormalc                 C   s`   t d}t d}tttjd|| tttjd|| tttjd|| tttjd|| d S )N      r   r	   r   r	   r,   )nparangeidentityassert_raises
ValueErrorr   pdfcdf)selfmucov r9   E/usr/lib/python3/dist-packages/scipy/stats/tests/test_multivariate.pytest_input_shape'   s   

z'TestMultivariateNormal.test_input_shapec                 C   s   t jd d\}}}t|||}t|jd t jd}t jd}t t jd}t|||}t|jd d\}}}t	|||}t|jd t jd}t jd}t t jd}t	|||}t|jd d S )N  )      ?g333333?      @r      )
r/   randomseedr   r4   r   ndimrandnabsr5   )r6   xmeanr8   r4   r5   r9   r9   r:   test_scalar_values/   s"   

z)TestMultivariateNormal.test_scalar_valuesc                 C   f   t jd t jd}t jd}t t jd}t|||}t|||}t|t 	| d S Nr<   r?   )
r/   r@   rA   rC   rD   r   logpdfr4   r   logr6   rE   rF   r8   d1d2r9   r9   r:   test_logpdfJ      z"TestMultivariateNormal.test_logpdfc                 C   l   t jd t jd}t|}t|}t|d d}t|d d}t|t | t|t | d S Nr<   r?   r	   )	r/   r@   rA   rC   r   rJ   r4   r   rK   r6   rE   rM   rN   Zd3Zd4r9   r9   r:   test_logpdf_default_valuesT      

z1TestMultivariateNormal.test_logpdf_default_valuesc                 C   rH   rI   )
r/   r@   rA   rC   rD   r   logcdfr5   r   rK   rL   r9   r9   r:   test_logcdfa   rP   z"TestMultivariateNormal.test_logcdfc                 C   rQ   rR   )	r/   r@   rA   rC   r   rV   r5   r   rK   rS   r9   r9   r:   test_logcdf_default_valuesk   rU   z1TestMultivariateNormal.test_logcdf_default_valuesc                 C   sl   t jd d}t j|}td|d D ]}t j||}t ||j}t||dd}t|j	j
| qd S )Nr<      r	   Tallow_singular)r/   r@   rA   rC   rangedotTr   r   cov_inforank)r6   nrF   Zexpected_ranksr8   Zdistnr9   r9   r:   	test_rankx   s   z TestMultivariateNormal.test_rankc              	   C   s~  dd }t ddD ]}tj|}t d|d D ]}tj||}t||j}t||f}||d |d |f< ||}t|t||j}	t||}
tt||dd}tt||dd}tt||	dd}t|j	j
| t|j	j
| t|j	j
| ||d | }||}||
}t|| t|| ||d | }||}||
}t|| t|| qq	d S )Nc                 S   s$   t j| | }tj|\}}}|S N)r/   r@   rC   scipylinalgsvd)ra   Murb   vr9   r9   r:   _sample_orthonormal_matrix   s   zXTestMultivariateNormal.test_degenerate_distributions.<locals>._sample_orthonormal_matrixr	   r?   TrZ   )r\   r/   r@   rC   r]   r^   zerosr   r   r_   r`   r4   r   rJ   )r6   rk   ra   rE   krb   Zcov_kkZcov_nnri   Zcov_rryZdistn_kkZdistn_nnZdistn_rrZpdf_kkZpdf_nnZpdf_rrZ	logpdf_kkZ	logpdf_nnZ	logpdf_rrr9   r9   r:   test_degenerate_distributions   sF   






z4TestMultivariateNormal.test_degenerate_distributionsc                 C   s   d}d}d}t || }|| }t j||ftd}t || d|| d | d f< ttj|d ttj|d |d |f t j	 t
t j|d |d |f d|f t|}t
|j| d S )Ng     @@d   r,   Zdtyper   r	   )r/   exprl   floatZfill_diagonalr   re   rf   detinfr   slogdetr   log_pdet)r6   Zlarge_total_logZnposZnzeroZlarge_entryra   r8   psdr9   r9   r:   test_large_pseudo_determinant   s   $z4TestMultivariateNormal.test_large_pseudo_determinantc                 C   s   t jd d}t j||}t ||j}t j|}t jdd|}t|||}t|||}t	dD ]3}t	dD ],}	t|||	f ||}
t
|
|||	f  t|||	f ||}
t
|
|||	f dd q<q6d S )Nr<   rY   r,   r+   gMbP?rtol)r/   r@   rA   rC   r]   r^   r   r4   r5   r\   r   )r6   ra   datar8   rF   XZdesired_pdfZdesired_cdfijactualr9   r9   r:   test_broadcasting   s    z(TestMultivariateNormal.test_broadcastingc                 C   sn   t ddd}d\}}|d }t|||}t|||}t|| t|||}t|||}t|| d S )Nr   r,   
   )333333??      ?)r/   linspacer   r4   r   r   r5   )r6   rE   rF   r8   scalerM   rN   r9   r9   r:   test_normal_1D   s   
z%TestMultivariateNormal.test_normal_1Dc                 C   s  t ddg}t ddgddgg}d}d|d  }t d	d|}t ||\}}t ||d
f}||d d d d d	f< ||d d d d df< t|||}	t|	|d	d}
t|	|dd}tj||d	 |d d d}tj||d |d d d}t	|
|ddd t	||ddd d S )Nr>   g      @r   皙?333333?i     r	   r   r,   Zaxisr   r   )locr   )r	   r	   {Gz?r{   atol)
r/   arrayr   Zmeshgridemptyr   r4   r&   r   r   )r6   rF   r8   ra   Zdeltarj   ZxvZyvposr4   Zmargin_xZmargin_yZgauss_xZgauss_yr9   r9   r:   test_marginalization   s    z+TestMultivariateNormal.test_marginalizationc                 C   s   t jd t jd}t jd}t t jd}t||}t||t||| t||t||| t|	|t	||| t|
|t
||| d S rI   )r/   r@   rA   rC   rD   r   r   r4   rJ   r5   rV   )r6   rE   rF   r8   Znorm_frozenr9   r9   r:   test_frozen  s   


z"TestMultivariateNormal.test_frozenc           	   	   C   s   t jd d}t j||}t ||j}tj|\}}t 	|d}d|d< d|d< t |t t 
||j}d}t||d	}t|j|d	}t|jt t |d d  t|j |j d S )
Nr<      r         ?r   Hz>h㈵>)cond)r/   r@   rA   rC   r]   r^   re   rf   Zeighfulldiagr   Zpinvr   rw   sumrK   )	r6   ra   rE   r8   rb   ri   r   rx   Zpsd_pinvr9   r9   r:   test_pseudodet_pinv  s    z*TestMultivariateNormal.test_pseudodet_pinvc                 C   s    g dg dg}t tt| d S )Nr	   r,   r+   )rY   r?   r   r2   r3   r   r6   r8   r9   r9   r:   test_exception_nonsquare_cov.     z3TestMultivariateNormal.test_exception_nonsquare_covc                 C   s@   ddgdt jgg}ttt| ddgdt jgg}ttt| d S )Nr	   r   )r/   nanr2   r3   r   ru   )r6   Zcov_nanZcov_infr9   r9   r:   test_exception_nonfinite_cov2  s   z3TestMultivariateNormal.test_exception_nonfinite_covc                 C   s    ddgddgg}t tt| d S )Nr	   r   r   r   r   r9   r9   r:   test_exception_non_psd_cov8  r   z1TestMultivariateNormal.test_exception_non_psd_covc                 C   s   t jd t jd}t jd}t d}t jj}t|t|| t|tj	||| t|tj
||| t|tj||| t|tj||| d S )Nr<   r?   )r?   r?   )r/   r@   rA   rC   onesrf   LinAlgErrorr2   r   r4   rJ   r5   rV   )r6   rE   rF   r8   er9   r9   r:   test_exception_singular_cov<  s   
z2TestMultivariateNormal.test_exception_singular_covc                 C   s"  t g d}t ddd}d| d }|t | }t |||gj}t g dd}t g dg d	g d
gd}t|||}t||dd t g d}	t|||}
t|
|	dd t g d}t ||gj}t ddgd}t ddgddggd}t|||}t||dd d S )N)g0Q-?gox'V?g[S, t?gDB?gc_.?r   r,   r?   r+   )r	   r+   r,   dr	   r,   r   )r,   r?   r   )r   r   r+   绽|=r   )gwk:E]?g̙Z?g0?gL8?gAG*?r   )gIى?gG?gƏ)?gėh?gEE?r	   )	r/   r   r   Zcosr^   r   r4   r   r5   )r6   Zr_pdfrE   rn   zrrF   r8   r4   Zr_cdfr5   Zr_cdf2Zr2mean2cov2Zcdf2r9   r9   r:   test_R_valuesH  s$   z$TestMultivariateNormal.test_R_valuesc                 C   s<   t d}t d}t||dd}| }t|ddg d S )Nr,   r,   r,   TrZ   r   )r/   rl   r   rvsr   )r6   rF   Z
covarianceZmodelsampler9   r9   r:   ,test_multivariate_normal_rvs_zero_covariance  s
   

zCTestMultivariateNormal.test_multivariate_normal_rvs_zero_covariancec                 C   s   d}d}t jt|d|d}t|j||f t jd tddgddgg|d}t|j|df t ddd}||}t|j|f d S )	Ni,  rY   r	   )rF   r8   sizer,   皙?r   rF   r8   )r   r   r/   rl   r   shaper   )r6   Nr   r   ri   r9   r9   r:   test_rvs_shape  s   
z%TestMultivariateNormal.test_rvs_shapec                 C   sx   t jd d}t j|}t j||}t ||j}d}t|||}tt	
|j|dd t|d|dd d S )N  r+   i  r   rz   r   )r/   r@   rA   rC   r]   r^   r   r   r   numpyr8   rF   )r6   ra   rF   rh   r8   r   r   r9   r9   r:   test_large_sample  s   z(TestMultivariateNormal.test_large_samplec                 C   s   t jd d}t j|}t j||}t ||j}t||}t| t|| t j	
|d }d|t dt j d  t t |  }t||  d S )Nr   r+   r   r   r,   r	   )r/   r@   rA   rC   r]   r^   r   r   entropyrf   ZeigrK   pir   )r6   ra   rF   rh   r8   rveigsdesiredr9   r9   r:   test_entropy  s   
,z#TestMultivariateNormal.test_entropyc                 C   s*   t g d}d}tt t|| d S )Nr	   r	   r	   r   )r/   r   r   rr   r   )r6   alphar   r9   r9   r:   test_lnB  s   zTestMultivariateNormal.test_lnBN)__name__
__module____qualname__r;   rG   rO   rT   rW   rX   rc   ro   ry   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r9   r9   r9   r:   r*   &   s2    

,@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 )TestMatrixNormalc              	   C   s  d}d}t ||fd}dt | t ||fd }dt | t ||fd }tttt d ttt|t d| ttt||t d ttt||| ttt||| ttt|j|| t jj	}t|t||t 
||f t|t|t 
||f| d S )NrY   r+   333333?r   ffffff?)r?   rY   r+   r   )r/   r   r1   r2   r3   r   rl   r^   rf   r   r   )r6   num_rowsnum_colsrh   UVr   r9   r9   r:   test_bad_input  s   zTestMatrixNormal.test_bad_inputc                 C   s  d}d}t ||fd}dt | t ||fd }dt | t ||fd }t ||f}t |df}t d|f}t |}	t |}
t d}ttj|||dj||f ttj|dj||f ttj|d	j|df ttj|d
jd|f ttj||dj||f ttj||dj||f ttj||dj||f tt|dj|	 tt|dj	|
 tt|d	j
| tt|d	j	| tt|d
j
| tt|d
j| tt||dj	|
 tt||dj|	 tt||dj
| d S )NrY   r+   r   r   r   r	   rF   rowcovcolcov)rF   )r   )r   )rF   r   )rF   r   )r   r   )r/   r   r1   rl   r   r   r   r   r   r   rF   )r6   r   r   rh   r   r   ZZZrZZcIrIcZI1r9   r9   r:   test_default_inputs  sR   


z$TestMatrixNormal.test_default_inputsc           
      C   s   d}d}t ||fd}t |d}d}t |d}d}t |}t |}	tt|||djd|  tt|||djd|	  tt|||djd|  tt|||djd|	  d S )NrY   r+   r   r   r   r   )r/   r   r1   r   r   r   r   )
r6   r   r   rh   ZUvZUsZVvZVsr   r   r9   r9   r:   test_covariance_expansion  s*   

z*TestMatrixNormal.test_covariance_expansionc                 C   s   t ddD ]w}t ddD ]o}t||fd}dt| t||fd }dt| t||fd }t|||d}|jdd}tj|||dd	}t|| |jdd}	||	}
tj|	|||d}t|
| ||	}tj|	|||d}t|| qqd S )
Nr	   r?   r   r   r   r   r<   random_state)rF   r   r   r   )	r\   r/   r   r1   r   r   r   r4   rJ   )r6   r~   r   rh   r   r   frozenrvs1rvs2r}   pdf1pdf2logpdf1logpdf2r9   r9   r:   test_frozen_matrix_normal  s(   




z*TestMatrixNormal.test_frozen_matrix_normalc                 C   s   t ddD ]u}t ddD ]m}t||fd}dt| t||fd }dt| t||fd }t|||d}|jdd}||}||}	|j	 }
|j	 }t
||}tj|
||d	}tj|
||d	}t||d
d t|	|d
d qqd S )Nr	   r?   r   r   r   r   r<   r   r   r   rz   )r\   r/   r   r1   r   r   r4   rJ   r^   ZflattenZkronr   r   )r6   r~   r   rh   r   r   r   r}   r   r   ZvecXZvecMr8   r   r   r9   r9   r:   test_matches_multivariate5  s$   



z*TestMatrixNormal.test_matches_multivariatec                 C   sD  d}d}t ||fd}dt | t ||fd }dt | t ||fd }d}t|||d}|j|dd	}|j|d
d	}	t j|t jd d d d d d f |	t jd d d d d d f fdd}
t|
jd|||f |	|
}t|jd|f t
dD ] }t
|D ]}tj	|
||f |||d}t||||f d qqd S )NrY   r+   r   r   r   r   r   r<   r   r   i  r   r   r,   r   )r/   r   r1   r   r   Zconcatenatenewaxisr   r   rJ   r\   r   )r6   r   r   rh   r   r   r   r   ZX1ZX2r}   Zarray_logpdfr~   r   Zseparate_logpdfr9   r9   r:   test_array_inputL  s*   F
z!TestMatrixNormal.test_array_inputc                 C   s   d}d}t ||fd}dt | t ||fd }dt | t ||fd }d}t|||d}|j|dd	}t j|d
d}	t|	|dd t ||| |j	}
t|
|dd t t 
|dd|| |j	}t||dd d S )NrY   r+   r   r   r     r   r<   r   r   r   r   r   r	   r,   )r/   r   r1   r   r   rF   r   r8   reshaper^   Zswapaxes)r6   r   r   rh   r   r   r   r   r}   Zsample_meanZsample_colcovZsample_rowcovr9   r9   r:   test_momentsc  s$   zTestMatrixNormal.test_momentsN)
r   r   r   r   r   r   r   r   r   r   r9   r9   r9   r:   r     s    'r   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'S )(TestDirichletc                 C   s   t jd t jdd}t jdd|}t|}t| t| t| t| t|	 t	| d}t
|D ]9}t jdd|}|t | }t||d d t|d d | t||d d t|d d | q=d S Nr   r	       &.>rp   r   r   )r/   r@   rA   randintr   r   r   varrF   r   r\   r   r4   rJ   r6   ra   r   r   	num_testsr~   rE   r9   r9   r:   test_frozen_dirichlet{  s   (*z#TestDirichlet.test_frozen_dirichletc                 C   s   t jd t g d}t jj|dd}t|jd tttj	|| tttj
|| t	|j| t	|jd d | t
|j| t
|jd d | d S )Nr   r          @      @r   r   )r   r+   r   )r/   r@   rA   r   r   r   r   r2   r3   r4   rJ   r^   r6   r   rE   r9   r9   r:   "test_numpy_rvs_shape_compatibility  s   z0TestDirichlet.test_numpy_rvs_shape_compatibilityc                 C   R   t jd g d}t jjt d|ddj}tttj|| tttj	|| d S )Nr   )r           r   r   r   r   
r/   r@   rA   r   Zmaximumr^   r2   r3   r4   rJ   r   r9   r9   r:   test_alpha_with_zeros  
   z#TestDirichlet.test_alpha_with_zerosc                 C   r   )Nr   )r   g       r   r   r   r   r  r   r9   r9   r:    test_alpha_with_negative_entries  r  z.TestDirichlet.test_alpha_with_negative_entriesc                 C   sp   t g d}t g d}t|| t|| t g d}tt||d tt||t d d S )Nr   r   r         @r   r   r   r   )r   r   r   r   r   )r/   r   r   r4   rJ   r   rK   r   r9   r9   r:   test_data_with_zeros  s   z"TestDirichlet.test_data_with_zerosc                 C   @   t g d}t g d}tttj|| tttj|| d S )N)r   r   r   r  r  r/   r   r2   r3   r   r4   rJ   r   r9   r9   r:   $test_data_with_zeros_and_small_alpha     z2TestDirichlet.test_data_with_zeros_and_small_alphac                 C   r	  )Nr  )r   皙r   r   r
  r   r9   r9   r:   test_data_with_negative_entries  r  z-TestDirichlet.test_data_with_negative_entriesc                 C   r	  )Nr  )r   皙?r   r   r
  r   r9   r9   r:    test_data_with_too_large_entries  r  z.TestDirichlet.test_data_with_too_large_entriesc                 C   >   t g d}t dd}tttj|| tttj|| d S )Nr   )r,   r   r   g$I$I?r/   r   r   r2   r3   r   r4   rJ   r   r9   r9   r:   test_data_too_deep_c     z"TestDirichlet.test_data_too_deep_cc                 C   sF   t ddgddgg}t dd}tttj|| tttj|| d S )Nr   r   r   r  )r,   r,   r         ?r  r   r9   r9   r:   test_alpha_too_deep  s   z!TestDirichlet.test_alpha_too_deepc                 C   s6   t g d}t dd}t|| t|| d S )Nr   r+   r   UUUUUU?)r/   r   r   r   r4   rJ   r   r9   r9   r:   test_alpha_correct_depth  s   z&TestDirichlet.test_alpha_correct_depthc                 C   r  )Nr   r  r   r  r   r9   r9   r:   test_non_simplex_data  r  z#TestDirichlet.test_non_simplex_datac                 C   r  )Nr  )r,   r   r   r  r   r9   r9   r:   test_data_vector_too_short  r  z(TestDirichlet.test_data_vector_too_shortc                 C   r  )Nr  )r?   r   r   r  r   r9   r9   r:   test_data_vector_too_long  r  z'TestDirichlet.test_data_vector_too_longc                 C   sF   t g d}t|}g d}g d}t| | t| | d S )N)r   皙?r   )gUUUUUU?g{Gz?gQ?r   皙?r   )r/   r   r   r   r   rF   )r6   r   r   Zexpected_varZexpected_meanr9   r9   r:   test_mean_and_var  s   zTestDirichlet.test_mean_and_varc                 C   s`   t dg}t|}t| jd t| jd t|dgjd t|dgjd d S )Nr   r   r   )	r/   r   r   r   rF   rB   r   r4   rJ   )r6   r   r   r9   r9   r:   rG     s   z TestDirichlet.test_scalar_valuesc                 C   s   t jd t jdd}t jdd|}t|}d}t|D ] }t jdd|}|t | }t|	|d d |	| qd S r   )
r/   r@   rA   r   r   r   r\   r   r   r4   r   r9   r9   r:    test_K_and_K_minus_1_calls_equal  s    z.TestDirichlet.test_K_and_K_minus_1_calls_equalc                 C   s   t jd t jdd}t jdd|}t|}d}d}d }t|D ]J}t|D ]}t jdd|}	|	t |	 }	|d urFt ||	f}q)|	}q)|	|j
}
d }|D ]}|	|}|d uret ||}qS|}qSt|
| q#d S )Nr   r	   r   r   rp   r   r?   )r/   r@   rA   r   r   r   r\   r   vstackr4   r^   appendr   )r6   ra   r   r   r   Znum_multipleZxmr~   mrE   ZrmZrsxsr   r9   r9   r:   test_multiple_entry_calls  s.   
z'TestDirichlet.test_multiple_entry_callsc                 C   s   t jd t jddd}t|}t|d |d }d}t|D ]}t jddd}|t | }t|	||	|g q!t|
 |
 d  t| | d  d S )Nr   r   rp   r,   r   r	   r   )r/   r@   rA   r   r   r   r\   r   r   r4   rF   r   )r6   r   r   br   r~   rE   r9   r9   r:   test_2D_dirichlet_is_beta)  s   z'TestDirichlet.test_2D_dirichlet_is_betaN)r   r   r   r   r   r  r  r  r  r  r  r  r  r  r  r  r  r   rG   r!  r&  r(  r9   r9   r9   r:   r   y  s(    	
r   c               
   C   s   t ddg} t dgg}ttt| | zt| | W d S  ty? } zd}tt|d t| | W Y d }~d S d }~ww )Nr   r   zDimension mismatch)r/   r   r2   r3   r   r   strlen)r7   Zsigmar   msgr9   r9   r:   ,test_multivariate_normal_dimensions_mismatch:  s   (r,  c                   @   4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )TestWishartc                 C   s  t jddd}ddgt dt jd t jdddg}|D ]}td|}t|j| t|jj|j qt ddgddgg}ddgt jd t ddgddggg}|D ]}td|}t|j| t|jj|j qRtttdt 	d tdt 	d t jddd}tttd| d S )Nr	   r,   Zndminr   )r	   r,   r  r+   )
r/   r   r_r   r   r   r   r2   r3   eye)r6   Z
true_scaleZscalesr   wr9   r9   r:   test_scale_dimensionsN  s8   



z!TestWishart.test_scale_dimensionsc              
   C   s  ddgt dt jd t jdddt jdgddg}tdd}|t jddd}|D ]
}t||| q+g dt jd t jg dddg}tdd}|t jg ddd}|D ]
}t||| qZdddgt dt jd t ddgddggt ddgddggd d d d t jf g}tdt d}|t ddgddggd d d d t jf }|D ]
}t||| qd S )Nr	   r,   r/  r+   r   r   r   )r/   r   r0  r   r4   r   r   r1  )r6   r}   r2  ZdensityrE   r9   r9   r:   test_quantile_dimensionsy  sR   
	


z$TestWishart.test_quantile_dimensionsc           	      C   sn  d}t t |d }t ||d  d |t j|dd< t |j|}g }tdD ].}t t ||d d  }t ||d  d |t j|dd< t |j|}|| q+t |j}ddt 	dddfd||fg}|D ]C\}}}t
||}t| t
|| t| t
|| t| t
|| t| t
|| t||t
||| qqd S 	NrY   r	   r,   r   rm   r?   r   r   )r/   r   r0   tril_indicesr]   r^   r\   r#  r   r   r   r   r   rF   moder   r4   )	r6   dimr   r}   r~   rE   
parametersdfr2  r9   r9   r:   r     s,   $$
zTestWishart.test_frozenc                 C   s   t jd d}d}t |}t jdddtd}t jdddd}|D ]C}t||}t|}t	|
 |
  t	| |  t	| |  t	|||| |j|d	}	|f}
d
}td|
||	 q"d S )N^   r	   r   r,   rq   r   Znumr   r   r   )r/   r@   rA   r1  r0   rs   r   r   r   r   r   rF   r   r4   r   r
   )r6   snr9  r   df_ranger}   r;  r2  cr   argsr   r9   r9   r:   test_1D_is_chisquared  s$   

z!TestWishart.test_1D_is_chisquaredc                 C   s  t jd d}d}d}t t dd }t d|t jddd< t |j|}t |df}|j||	 }t
||}t||d	}t| |  t| |  t| |  t jd
ddd}	t||	||	 |j|d}
|d|f}d}td|||
 d S )Nr<  r=  r   rY   r	   r   r   r6  r   r   r>  r   r   r   r   )r/   r@   rA   r   r0   r7  r]   r^   r   Zsqueezer   r   r   r   rF   r   r   r4   r   r
   )r6   r?  r;  r9  r   ZlamdaZsigma_lamdar2  rA  r}   r   rB  r   r9   r9   r:   test_is_scaled_chisquared  s(   

z%TestWishart.test_is_scaled_chisquaredN)r   r   r   r3  r4  r   rC  rE  r9   r9   r9   r:   r.  M  s    +1 r.  c                   @   st   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S )TestMultinomialc                 C   sl   t ddd}t|ddd t ddgd	d
dg}t|tjdd t ddgd	ddg}t|tjdd d S )Nr+   rY   r   r   r   g&Mty:0yE>rz   r+   rY   r   r   r   )r   logpmfr   r/   ZNAN)r6   vals1vals2vals3r9   r9   r:   test_logpmf  s   zTestMultinomial.test_logpmfc                 C   sX   t ddd}tddd}t||dd t dd	d
}tdd	d}t||dd d S )NrG  r   rH  r+   r   rI  rz   )r         r   r   r   r   )r   rK  r!   r   pmfr6   val1val2r9   r9   r:   test_reduces_binomial  s   z%TestMultinomial.test_reduces_binomialc                 C   sR   dg d}}ddddddd	dd
dd
}|D ]}t t||||| dd qd S )Nr+   )g      ?r  g      ?g     @?g     ?g      ?g      `?g     ?g      ?g      ?g      ?g      ?)
)r   r   r+   )r	   r   r,   )r,   r   r	   )r+   r   r   r.   r   )r,   r	   r   )r   r,   r	   r   )r   r+   r   g+=r   )r   r   rS  )r6   ra   pZr_valsrE   r9   r9   r:   test_R"  s   zTestMultinomial.test_Rc                 C   sH   t jddgd ddd}tjd}|j ddgd dd}t|| d S )Nr+   r  r   {   r   r   )r   r   r/   r@   RandomStater   )r6   Zsc_rvsrndmZnp_rvsr9   r9   r:   test_rvs_np3  s   zTestMultinomial.test_rvs_npc                 C   s   t ddd}t|ddd t ddd	}t|d
dd t ddgddggddgddgggdd}t|ddgddggdd tjdtjd}t |dd	}t|tjg tjd t ddgdd	}t|ddd t g ddg d}t|ddd d S )N)r?   r?   r	   r	   rI  rz   rG  r   rH  g5
?r+   r   rP  r   	   rR  g*7?g<Wƌ?r   r,   rq   rY   r,   r+   r+   r   r   )gUUUUUU?r  r   gߦ?)r   rS  r   r/   r   float64r   )r6   vals0rL  rM  rE   rN  vals4Zvals5r9   r9   r:   test_pmf:  s   $zTestMultinomial.test_pmfc                 C   s
  t ddgdddgddgg}t|dd	gd
d t ddgddgddg}t|ddgd
d t ddgddgggdddg}t|ddggd
d t ddgdgdgggddg}t|dgdgggd
d t ddgddggdggggddg}t|ddggggd
d d S )Nr	   r,   r+   r   r   r   r  gv?g~jt?rI  rz   rY   r   )r   rS  r   )r6   rc  rL  rM  rN  rd  r9   r9   r:   test_pmf_broadcastingO  s     &z%TestMultinomial.test_pmf_broadcastingc                 C   s4   t dd}g dg dg dg}t||dd d S )Nr?   )r   r   r   )r  g333333ӿ      )g433333ӿg?      )rg  rh  g      ?rI  rz   r   r8   r   r6   cov1r   r9   r9   r:   test_cov_  s   zTestMultinomial.test_covc                 C   s   t dddgddgg}ddgddggddgddggg}t||d	d
 t ddgddg}ddgddggddgddggg}t||d	d
 t ddgddgddgg}ddgddggddgddggg}t||d	d
 d S )Nr?   r   r   r   r  g?gܿgrI  rz   rY   g
ףp=
?g
ףp=
׿r   r   r  r   gzG?gzGr   g333333ri  r6   rk  r   cov3cov4cov5cov6r9   r9   r:   test_cov_broadcastingf  s     z%TestMultinomial.test_cov_broadcastingc                 C   s*   t dddg}t|tdddd d S )Nr,   r   r  rI  rz   r   r   r   r!   )r6   ent0r9   r9   r:   r   t  s   zTestMultinomial.test_entropyc                 C   s   t ddgddg}t|tddtddgdd t ddgdd	gd
dgg}t|tddtdd
gdd t dgdggdd	gd
dgg}t|tddtdd
gtddtdd
ggdd d S )Nr,   r+   r   r   rI  rz   r   rP  r   r  r   rs  )r6   rt  Zent1Zent2r9   r9   r:   test_entropy_broadcastingz  s    
z)TestMultinomial.test_entropy_broadcastingc                 C   s&   t dddg}t|ddgdd d S )Nr?   r   r  r   r  rI  rz   r   rF   r   r6   mean1r9   r9   r:   	test_mean  s   zTestMultinomial.test_meanc                 C   s2   t ddgddg}t|ddgddggd	d
 d S )Nr?   r   r   r  r   r  g433333?g433333@rI  rz   rv  rw  r9   r9   r:   test_mean_broadcasting  s   z&TestMultinomial.test_mean_broadcastingc                 C   s   t jd d}d}g dg dg dg dg dg}t j|t jd	}t||}t||t||| t||t||| t|	 t	|| d S )
Nr<      )r   r   r   r  r   r   r   r{  r   r   r	      r   r	   r	   r   r	   r	   r	   r_  r	   r	   r,   rP  rq   )
r/   r@   rA   asarrayrb  r   r   rS  rK  r   )r6   ra   ZpvalsrE   Z	mn_frozenr9   r9   r:   r     s   "
zTestMultinomial.test_frozenN)r   r   r   rO  rW  rY  r]  re  rf  rl  rr  r   ru  ry  rz  r   r9   r9   r9   r:   rF    s    
rF  c                   @   r-  )TestInvwishartc           	      C   sX  d}t t |d }t ||d  d |t j|dd< t |j|}g }tdD ].}t t ||d d  }t ||d  d |t j|dd< t |j|}|| q+t |j}ddt 	dddfd||fg}|D ]8\}}}t
||}t| t
|| t| t
|| t| t
|| t||t
||| qqd S r5  )r/   r   r0   r7  r]   r^   r\   r#  r   r   r   r   r   rF   r8  r   r4   )	r6   r9  r   r}   r~   rE   r:  r;  iwr9   r9   r:   r     s*   $$
zTestInvwishart.test_frozenc                 C   s   t jd d}d}t |}t jdddtd}t jdd	d	d
}|D ]B}t||}t|d dd}t	|
 |
  t	| |  t	|||| |j|d}	|d ddf}
d}td|
||	 q"d S )Nr<  r=  r	   r?      r,   rq   r   r   r>  r   rD  r   r   r   r   )r/   r@   rA   r1  r0   rs   r   r   r   r   r   rF   r4   r   r
   )r6   r?  r9  r   r@  r}   r;  r  Zigr   rB  r   r9   r9   r:   test_1D_is_invgamma  s"   

z"TestInvwishart.test_1D_is_invgammac                 C   s|  d}d}t |}d|d< d|d< t||}t||}t jd t||}t jd | }t jd t||}t jd | }	t jd t jjdd}
t jt j	|t j	|d t j	|d	 f d }t 
|}|
|t j|d
d< t j|}||}t ||j}t jt j|}||}t jt ||j}t|| t|| t|| t|	| d S )Nr+   r   r   r-   r	   r   i r   r	   r,   r   r6  )r/   r1  r   r   r@   rA   r   normalr0  Z	chisquarer   r7  rf   Zcholeskyr]   r^   invr   )r6   r9  r;  r   r2  r  Zw_rvsZfrozen_w_rvsZiw_rvsZfrozen_iw_rvsZcovariancesZ	variancesADZDAZmanual_w_rvsZiDZiDAZmanual_iw_rvsr9   r9   r:   test_wishart_invwishart_2D_rvs  sJ   









z-TestInvwishart.test_wishart_invwishart_2D_rvsc                 C   s   t g dg dg dg dg}t g dg dg dg dg}t ||g}| }t| t d}t|d	 |d	 |d
d t|d |d |d
d dS )Regression test for gh-8844.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   r	   rY   rY   r   V瞯<r   r	   N)r/   r   copyr   r1  r   r]   )r6   Za0Za1aZainvZidentr9   r9   r:   test_cho_inv_batch  s    


 z!TestInvwishart.test_cho_inv_batchc                 C   s   t g dg dg dg dg}t g dg dg dg dg}d	}t|||}|jd
 }t j|\}}t j|\}}t j||}	|d | || d t d  t	|d | || d d |  d|	
   }
t||
 dS )r  r  r  r  r  )r_  r   r+   r	   )r   r_  r?   r	   )r+   r?   rP  r,   )r	   r	   r,   r_  r   r   r,   r	   r   N)r/   r   r   rJ   r   rf   rv   ZsolverK   r'   Ztracer   )r6   r}   ZPsiZnuZprobrX  ZsigZlogdetXZ	logdetPsirh   expectedr9   r9   r:   test_logpdf_4x4(  s4   




zTestInvwishart.test_logpdf_4x4N)r   r   r   r   r  r  r  r  r9   r9   r9   r:   r    s     ;r  c                   @   r-  )TestSpecialOrthoGroupc                 C   sf   t jd td}t g dg dg dg}t|| t jjdd}tjd|d}t|| d S )N  r+   )g(fg^A.g*B?)g:d?g2ugi}?)gR#xlF?g^`D?g?rA   r   )r/   r@   rA   r   r   r   r   r[  )r6   rE   r  r   r9   r9   r:   test_reproducibilityB  s   


z*TestSpecialOrthoGroup.test_reproducibilityc                 C   <   t ttjd  t ttjd t ttjd t ttjd d S Nr   r	   r>   )r2   r3   r   r   r6   r9   r9   r:   test_invalid_dimN     z&TestSpecialOrthoGroup.test_invalid_dimc                 C   s4   d}t |}|jdd}t j|dd}t|| d S )Nr   r<   r   )r   r   r   )r6   r9  r   r   r   r9   r9   r:   test_frozen_matrixT  s
   z(TestSpecialOrthoGroup.test_frozen_matrixc                 C   sd   dd t ddD }dd |D }t|dgd dd	 |D ]}tt||jt|jd
  qd S )Nc                 S   $   g | ]}t d D ]}t|qqS r+   )r\   r   r   .0r9  r~   r9   r9   r:   
<listcomp>^      z<TestSpecialOrthoGroup.test_det_and_ortho.<locals>.<listcomp>r,   r{  c                 S      g | ]}t j|qS r9   r/   rf   rt   r  rE   r9   r9   r:   r  c      r      vIh%<=rz   r   )r\   r   r   r/   r]   r^   r1  r   )r6   r%  detsrE   r9   r9   r:   test_det_and_ortho]  s   z(TestSpecialOrthoGroup.test_det_and_orthoc                    |   d}d}d}t jd tj||dd tfdd D  fd	d
 D }fdd
|D }t|gt| | d S )Nr?   r   皙?r  r   r   r`  )r	   rY   )r,   r+   c                 3   4    | ]\  ft  fd dD fV  qdS )c                       g | ]}|   qS r9   r9   r  ZecZerr9   r:   r    r  z=TestSpecialOrthoGroup.test_haar.<locals>.<genexpr>.<listcomp>Nsortedr  r%  r  r:   	<genexpr>     2 z2TestSpecialOrthoGroup.test_haar.<locals>.<genexpr>c                    &   g | ]} D ]
}||kr||fqqS r9   r9   r  Ze0Ze1elsr9   r:   r       & z3TestSpecialOrthoGroup.test_haar.<locals>.<listcomp>c                    &   g | ]\}}t  |  | d  qS r^  r   r  Zp0Zp1projr9   r:   r    r  )r/   r@   rA   r   r   dictr   r*  r6   r9  samplesZks_probZpairsZks_testsr9   r  r  r%  r:   	test_haark     
zTestSpecialOrthoGroup.test_haarN)r   r   r   r  r  r  r  r  r9   r9   r9   r:   r  A  s    	r  c                   @   s<   e Zd Zdd Zdd Zdd Zdd Zejj	d	d
 Z
dS )TestOrthoGroupc                 C   sj   t jd td}tjddd}tt j|d t g dg dg dg}t	|| t	|| d S )Ni  r+   r   r   )gmS9?g8e,˿gRVϿ)gob2g	,ʶgC}Ȇ?)gF-YD?g:kO?g-C?)
r/   r@   rA   r   r   r   rf   rt   r   r   r6   rE   x2r  r9   r9   r:   r    s   


z#TestOrthoGroup.test_reproducibilityc                 C   r  r  )r2   r3   r   r   r  r9   r9   r:   r    r  zTestOrthoGroup.test_invalid_dimc                 C   s   dd t ddD }tdd |D }tt|t|jdd tdgd	 d
d |D  tdd |D d	gd	  |D ]}|D ]}tt	||j
t|jd  qBq>d S )Nc                    s"   g | ]  fd dt dD qS )c                    s   g | ]}t  qS r9   )r   r   r  r~   r9  r9   r:   r    s    @TestOrthoGroup.test_det_and_ortho.<locals>.<listcomp>.<listcomp>r   )r\   r  r9   r  r:   r    s
    

z5TestOrthoGroup.test_det_and_ortho.<locals>.<listcomp>r,   r{  c                 S   s   g | ]	}d d |D qS )c                 S   r  r9   r  r  r9   r9   r:   r    r  r  r9   )r  xxr9   r9   r:   r        r  rz   r   r   c                 S   $   g | ]}t |d k d  jd  qS r   r/   Znonzeror   r  r   r9   r9   r:   r       $ c                 S   r  r  r  r  r9   r9   r:   r    r  )r\   r/   r   r   fabsr   r   r   r   r]   r^   r1  )r6   r%  r  r  rE   r9   r9   r:   r    s   z!TestOrthoGroup.test_det_and_orthoc                    r  )Nr?   r   r  i  r   r  c                 3   r  )c                    r  r9   r9   r  r  r9   r:   r    r  z6TestOrthoGroup.test_haar.<locals>.<genexpr>.<listcomp>Nr  r  r  r  r:   r    r  z+TestOrthoGroup.test_haar.<locals>.<genexpr>c                    r  r9   r9   r  r  r9   r:   r    r  z,TestOrthoGroup.test_haar.<locals>.<listcomp>c                    r  r^  r  r  r  r9   r:   r    r  )r/   r@   rA   r   r   r  r   r*  r  r9   r  r:   r    r  zTestOrthoGroup.test_haarc                    sh   t jd dd }tddD ]" d fdd		}||}|tjjj}tj||\}}t	d
| qd S )Nr  c                 S   s,   t jt jj| | fd\}}}t ||S Nr   )r/   rf   rg   r@   r  r]   )r9  ri   Z_srj   r9   r9   r:   random_ortho  s    z<TestOrthoGroup.test_pairwise_distances.<locals>.random_orthor,   r   r   r   c                    s<   t  fddt|D }|t jj| ||jd7 }|S )Nc                    s*   g | ]}t  d  d  d qS )r  r,   )r/   r   )r  _)r9  r   r9   r:   r    s    z\TestOrthoGroup.test_pairwise_distances.<locals>.generate_test_statistics.<locals>.<listcomp>r   )r/   r   r\   r@   r   r   )r   r   Zepsstatsr  )r   r:   generate_test_statistics  s
   zHTestOrthoGroup.test_pairwise_distances.<locals>.generate_test_statisticsr  )r   r   )
r/   r@   rA   r\   re   r  r   r   r   r   )r6   r  r  r  r   Z_DrX  r9   r  r:   test_pairwise_distances  s   	z&TestOrthoGroup.test_pairwise_distancesN)r   r   r   r  r  r  r  pytestmarkZslowr  r9   r9   r9   r:   r    s    r  c                   @   ,   e Zd Zdd Zdd Zdd Zdd Zd	S )
TestRandomCorrelationc                 C   sb   t jd d}t|}tj|dd}t g dg dg dg dg}t|| t|| d S )Nr  )r   r  r   r=   r   )r   ʿ3LU?iƖ\)r  r   ``+ϿO0)?)r  r  r   qɨ?)r  r  r  r   )r/   r@   rA   r   r   r   r   )r6   r   rE   r  r  r9   r9   r:   r    s   


z*TestRandomCorrelation.test_reproducibilityc                 C   s|   t ttjd  t ttjd t ttjd t ttjdg t ttjddgddgg t ttjddg t ttjg d d S )	Ntestr>   r	   r,   r+   rY   rg  )r	   r,   r   )r2   r3   r   r   r  r9   r9   r:   test_invalid_eigs  s   z'TestRandomCorrelation.test_invalid_eigsc           
         s   dd  t jd  fddtddD }|g d d	d |D }d
d |D }dd |D }dd |D }t||ddd dd |D }t||D ]\}}t||dd qM|D ]
}	t|	|	jdd q[d S )Nc                 S   s   | | t | S rd   )r   )r~   r   r9   r9   r:   r     s   z3TestRandomCorrelation.test_definition.<locals>.normrZ  c                    s    g | ]} |t jj|d qS )r   )r/   r@   r   r  r   r9   r:   r    s     z9TestRandomCorrelation.test_definition.<locals>.<listcomp>r,   r   )rY   r   r   r   c                 S   s   g | ]	}d gt | qS )r   )r*  r  r   r9   r9   r:   r  	  r  c                 S      g | ]}t |qS r9   )r   r   r  r9   r9   r:   r  
      c                 S   s   g | ]}t t j|qS r9   )r/   r  rf   rt   r  r9   r9   r:   r    s    c                 S   r  r9   )r/   Zprodr  r9   r9   r:   r    r  r  r   c                 S   r  r9   )r/   r   r  r9   r9   r:   r    r  rz   )r/   r@   rA   r\   r#  r   zipr^   )
r6   r   r   r%  r  Z
dets_knownZdiagsr  r'  rE   r9   r  r:   test_definition  s   z%TestRandomCorrelation.test_definitionc                 C   sz  t jddgddggtd}t|}t|t ddgddgg t jddW t ddgddgg}t jddgdt ddggtd}t| }t||j	
|
| t jd	dgdd
ggtd}t| }t||j	
|
| W d    n1 sw   Y  t jddgddggtd}t| }t|d d t jddgddggtd}t| }t|d d d S )Nr   r   r	   rq   ignore)Zoverr   gu <7~r   r  r,   r   gMk   @)r/   r   rs   r   Z_to_corrr   ZerrstateZ	nextafterr  r^   r]   )r6   r$  gZm0r9   r9   r:   test_to_corr  s$   
"z"TestRandomCorrelation.test_to_corrN)r   r   r   r  r  r  r  r9   r9   r9   r:   r    s
    	$r  c                   @   r  )
TestUnitaryGroupc                 C   sX   t jd td}tjddd}t g dg dg dg}t|| t|| d S )Nr  r+   r   )y~p?yZ?y뉦?ĭ?y8Z?&4?)yEHξr?8i?ye??+ݓݿyu=ٿB۽'?)yqXQÿ?yN6ҿzܷy! 8l?2W'?)r/   r@   rA   r   r   r   r   r  r9   r9   r:   r  ?  s   


z%TestUnitaryGroup.test_reproducibilityc                 C   r  r  )r2   r3   r   r   r  r9   r9   r:   r  K  r  z!TestUnitaryGroup.test_invalid_dimc                 C   sJ   dd t ddD }|D ]}tt|| jt|jd dd qd S )Nc                 S   r  r  )r\   r   r   r  r9   r9   r:   r  R  r  z3TestUnitaryGroup.test_unitarity.<locals>.<listcomp>r,   r{  r   r  r   )r\   r   r/   r]   Zconjr^   r1  r   )r6   r%  rE   r9   r9   r:   test_unitarityQ  s   *zTestUnitaryGroup.test_unitarityc                 C   sx   d}d}t jd tj||d}t dd |D }t |j|j}t	|
 tt j dt j j}t|jdk d S )	Nr?   r   r  r   c                 S   r  r9   )re   rf   Zeigvalsr  r9   r9   r:   r  f  r  z.TestUnitaryGroup.test_haar.<locals>.<listcomp>r,   r  )r/   r@   rA   r   r   r"  Zarctan2imagrealr    Zravelr   r   r5   r   Zpvalue)r6   r9  r  r%  r   rE   resr9   r9   r:   r  Z  s    zTestUnitaryGroup.test_haarN)r   r   r   r  r  r  r  r9   r9   r9   r:   r  >  s
    	r  c                   @   s  e Zd Zddgddgddgddgddgddgddgddgddgddgg
ddgddgddggdg dfg dg d	g d
g dg dgg dg dg dg dgdg dfgZejdedd Zejdedd Zdd Z	dd Z
eddd Zdd  Zd!d" Zd#d$d%d&dgdd&ggd%ddgd&dgdd&ggdfd%d&dgdd&ggd&ddgd&dgdd&ggd&fd&d&gd%d%d&d&gddgddggdfd&d&gd%d&d&d&gddgddggd&fd&d&gd&dgdd&ggd%d&d&gd&dgdd&ggdfd&d&gd&dgdd&ggd&d&d&gd&dgdd&ggd&fgZejd'ed(d) Zd*ddd*gdggdfd*gdgdd*gdggdfed*gedgdd*gdggdfgZejd'ed+d, Zd-d. Zd/d0 Zd1d2 Zejd3d4d5gejd6g d7ejd8d9d:ejgd;d< Zd%S )=TestMultivariateTr	   r,   rY   r+   r?   r   )
|	ߌ?v
2R?r  g0$H?r  r  gO`?gՒXO?g	ߌ,?g8N8?)gHP?gQI?go_?)g9m4?g=U?gˡE?)g$C?ghs?gǺ?)g!uq?g%C?gS㥛?)gOec?gz,C?g	h"l?)r   r	   2   )r   r   r  )r   r   r  )r  r  r   rP  )gJ#R<g^<gL[<gN=<g`$<zx, loc, shape, df, ansc                 C   s(   t |||dd}||}t|| d S Nr   r  )r#   r4   r   )r6   rE   r   r   r;  ansdistvalr9   r9   r:   test_pdf_correctness  s   
z&TestMultivariateT.test_pdf_correctnessc           	      C   s8   t |||dd}||}||}tt|| d S r  )r#   r4   rJ   r   r/   rK   )	r6   rE   r   r   r;  r  r   rU  rV  r9   r9   r:   test_logpdf_correct  s   

z%TestMultivariateT.test_logpdf_correctc                 C   s.   g d}t j|dd}t|}t|| d S )N)
r_  r   rY   r	   r_  r   r  r   r+   r	   r;  )r#   r4   r$   r   )r6   rE   r  r  r9   r9   r:   test_mvt_with_df_one_is_cauchy  s   
z0TestMultivariateT.test_mvt_with_df_one_is_cauchyc                 C   s   d}t ddddd}|jdd}t|\}}||ksJ t ddgd	d
gd
d	ggddd}|jdd}t|\}}||k s@J d S )Nr   r   r	   i r;  rA   r   rJ  r+   r   r   *   )r#   r   r%   all)r6   Z	P_VAL_MINr   r  r  rX  r9   r9   r:   &test_mvt_with_high_df_is_approx_normal  s   z8TestMultivariateT.test_mvt_with_high_df_is_approx_normalz'scipy.stats.multivariate_normal._logpdfc                 C   s`   t ddtjdd}t|tsJ t jdtjd |jdksJ t jdtjd |jdks.J d S )Nr   r	   r   r  r  r,   )r#   r/   ru   
isinstancer   r4   Z
call_countrJ   )r6   Zmockr   r9   r9   r:   !test_mvt_with_inf_df_calls_normal  s   z3TestMultivariateT.test_mvt_with_inf_df_calls_normalc           	      C   s   d}t |}t |}d}t |}t||||}t |s#J t||||}t |s3J d}t j||f}t||||}|j|fksNJ t||||}|j|fks_J tt dt dd	 }t |stJ d}tt dt ddj	|d}|j|fksJ d S )NrY         @r   r	   r   )
r/   rl   r1  r#   r4   ZisscalarrJ   r@   r   r   )	r6   r9  r   r   r;  rE   r  Z	n_samplesr   r9   r9   r:   test_shape_correctness  s(   


 z(TestMultivariateT.test_shape_correctnessc                 C   s6   t  }t|jdg t|jdgg |jdksJ d S )Nr   r	   r#   r   r   r   r;  )r6   r   r9   r9   r:   test_default_arguments  s   z(TestMultivariateT.test_default_arguments)NNNr   r	   r	   )NNr   r   r	   r   Nr   z*loc, shape, df, loc_ans, shape_ans, df_ansc                 C   s8   t |||d}t|j| t|j| |j|ksJ d S )Nr   r   r;  r  r6   r   r   r;  Zloc_ansZ	shape_ansZdf_ansr   r9   r9   r:   test_default_args  s   z#TestMultivariateT.test_default_argsr   c                 C   s4   t |||}t|j| t|j| t|j| d S rd   r  r  r9   r9   r:   &test_scalar_list_and_ndarray_arguments"  s   z8TestMultivariateT.test_scalar_list_and_ndarray_argumentsc                 C   s   ddgg}t ttfi t|d ddgddgddgg}t ttfi t||d td}td}d}t ttfi t|||d d}t ttfi t|||d d S )	Nr	   )r   r,   r+   )r   r   r   r  r   )r2   r3   r#   r  r/   rl   r1  )r6   r   r   r;  r9   r9   r:   test_argument_error_handling)  s4   




z.TestMultivariateT.test_argument_error_handlingc                 C   sh   t jd}|jdd}t d}t||ddd}t||ddd}|jdd}|jdd}t|| d S )NrY   r+   r   r,   r  r   )r/   r@   r[  r   r1  r#   r   r   )r6   rngr   r   Zdist1Zdist2Zsamples1Zsamples2r9   r9   r:   r  B  s   
z&TestMultivariateT.test_reproducibilityc                 C   s:   t ddgddgddggddd}ttjjtfi | d S )Nr   r	   F)r   r   r;  r[   )r  r2   r/   rf   r   r#   )r6   rB  r9   r9   r:   test_allow_singularL  s    z%TestMultivariateT.test_allow_singularr   )r   r+   )r?   r   rY   r+   r9  )r,   r+   rY   r?   r;  r   r   c                 C   s<   t t|t||}|j|d}|j||f ksJ d S r  )r#   r/   rl   r1  r   r   )r6   r   r9  r;  r   r   r9   r9   r:   test_rvsQ  s   zTestMultivariateT.test_rvs)r   r   r   Z	PDF_TESTSr  r  parametrizer  r  r  r
  r)   r  r  r  ZDEFAULT_ARGS_TESTSr  r/   r   ZARGS_SHAPES_TESTSr  r  r  r  ru   r  r9   r9   r9   r:   r  l  s    
(B


 **""..
$

r  c                   @   s  e Zd Zejdddgddgddfddgddgdejfd	dgddgdejfddgd
dgdejfddgddggddgd
dggddgejejgfd	dgd
dgdejfddgddgdejfddgddgdejfddgddgdejfddgddgdejfg
dd Z	dd Z
dd Zdd Zejddgdgddfddgddgddfddgddggddgddgggddgddgddggdd gdd!ggfeg ejeg ejdg fddgddgddfg d"g d#d$d%fgd&d' Zejdddgddgdd(ggddd)gfdgdggdgdggddgd*d+gfdgdgggdgdggddgd*d+ggfdgdggdggggddgd*d+gggfg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dDS )ETestMultivariateHypergeomzx, m, n, expectedr+   rY   r?   r   r   gr   r  r	   r,   iir~  r{  r   ic                 C       t |||}t||dd d S )Ngư>rz   )r   rK  r   r6   rE   r$  ra   r  Zvalsr9   r9   r:   rO  [  s   z%TestMultivariateHypergeom.test_logpmfc                 C   st   t jddgddgdd}tjddddd}t||d	d
 t jddgddgdd}tjddddd}t||d	d
 d S )Nr+   r	   r   r?   rY   )rE   r$  ra      )rm   rh   ra   r   rI  rz   r      )r   rS  r"   r   rT  r9   r9   r:   test_reduces_hypergeomy  s   z0TestMultivariateHypergeom.test_reduces_hypergeomc                 C   s:   t ddgdd}|jddd}t|d| d	d
 d S )Nr+   r?   rY   r$  ra   r   rZ  r   r   r   rz   r   r   r   rF   r6   r   r   r9   r9   r:   r    s   z"TestMultivariateHypergeom.test_rvsc                 C   sF   t ddgddggddgd}|jddd	}t|d
| dd d S )Nr+   r?   r   rY   r_  r"  )r   r,   rZ  r   r   r   rz   r#  r$  r9   r9   r:   test_rvs_broadcasting  s   z/TestMultivariateHypergeom.test_rvs_broadcastingg zl?rP  r_  gQ?gE|?g
y?ra  )r?   r   r   r   gzp?c                 C   r  Nr   rz   r   rS  r   r  r9   r9   r:   re    s   z"TestMultivariateHypergeom.test_pmfr  g ?r   r   c                 C   r  r&  r'  r  r9   r9   r:   rf    s   
z/TestMultivariateHypergeom.test_pmf_broadcastingc                 C   s:   t jg ddd}g dg dg dg}t||dd d S )	N)r+   r   r   r{  r"  )gf_?&BUп疈Cx@ؿ)r(  ggs>d?ZsK)r)  r*  g35?rI  rz   r   r8   r   rj  r9   r9   r:   rl    s   z"TestMultivariateHypergeom.test_covc                 C   s   t jddgddggddgd}dd	gd	dggd
dgdd
ggg}t||dd t jdgdggddgd}dggdggg}t||dd t jddgddgd}dd	gd	dggddgddggg}t||dd d S )Nr   r_  r   r  rP  r{  r"  g?gg(\?g(\rI  rz   rY   r?   r   g333333?g333333r+  rm  r9   r9   r:   rr    s   z/TestMultivariateHypergeom.test_cov_broadcastingc                 C   s4   t jddgdd}tjdddd}t||dd d S )	Nr   r?   rY   r"  r  rh   ra   r   rI  rz   )r   r   r"   r   )r6   var0var1r9   r9   r:   test_var  s   z"TestMultivariateHypergeom.test_varc                 C   s   t jddgddgd}t jddgdd}t jddgdd}t|d |dd t|d	 |dd t jddgdd
ggddgd}ddgddgg}t||dd t jdgdggddgd}dgdgg}t||dd d S )Nr   r?   rY   rP  r"  r   rI  rz   r	   rQ  g̍3eY?gI{?r   )r   r   r   )r6   r-  r.  var2var3Zvar4Zvar5Zvar6r9   r9   r:   test_var_broadcasting  s   z/TestMultivariateHypergeom.test_var_broadcastingc                 C   s`   t jddgdd}tjdddd}t|d |dd	 t jd
dgdd}ddg}t||dd	 d S )Nr   r?   rY   r"  r  r,  r   rI  rz   r{  rP  g      @r  )r   rF   r"   r   r6   mean0rx  r   Zmean3r9   r9   r:   ry    s   z#TestMultivariateHypergeom.test_meanc                 C   s@   t jddgddggddgd}ddgd	d
gg}t||dd d S )Nr+   r?   r   rY   rP  r"  r=   r>   gUUUUUU@gUUUUUU@rI  rz   )r   rF   r   )r6   r4  rx  r9   r9   r:   rz    s   z0TestMultivariateHypergeom.test_mean_broadcastingc                 C   s   t jg ddd}t|g d t jg ddd}t|tjtjtjg t jg dg dgdd}t|tjtjtjgg dgd	d
 t jtg tjdd}t|g  t|j	dk d S )Nr   r   r   r   r"  r   r   r   r	   r   r   r,   r	   r   r	   )r   r   r   FFg<rz   r  )
r   rF   r   r/   r   r   r   int_r   r   r3  r9   r9   r:   test_mean_edge_cases  s   
z.TestMultivariateHypergeom.test_mean_edge_casesc                 C   s   t jg ddd}t|g ddd t jg ddd}t|tjtjtjg t jg dg d	gdd}t|tjtjtjgg dgd
d t jtg tjdd}t|g  t|j	dk d S )Nr5  r   r"  r6  gؗҜ<rz   r7  r,   r8  r9  r  )
r   r   r   r   r/   r   r   r:  r   r   )r6   r-  r.  r0  r1  r9   r9   r:   test_var_edge_cases  s   
z-TestMultivariateHypergeom.test_var_edge_casesc                 C   s   t jg ddd}g dg dg dg}t||dd t jg ddd}g dg dg dg}t|| t jtg tjdd}tjg tjd	dd}t||dd t	|j
d
k d S )Nr7  r	   r"  r6  r9  rz   r5  r   rq   r   )r   r8   r   r   r/   r   r:  Zfloat_r   r   r   )r6   Zcov0rk  rn  ro  rp  rq  r9   r9   r:   test_cov_edge_cases	  s   
z-TestMultivariateHypergeom.test_cov_edge_casesc                 C   s   t jd d}g d}g dg dg dg dg dg}t j|t jd	}t||}t||t||| t||t||| t|	 t	|| t|
 t
|| d S )
Nr<   r{  )r   r_  r~     r|  r}  r  r  r  rq   )r/   r@   rA   r  r:  r   r   rS  rK  r   r8   )r6   ra   r$  rE   Z
mhg_frozenr9   r9   r:   r     s    


z%TestMultivariateHypergeom.test_frozenc                 C   s   t ttjddd t ttjddgd t ttjddgdgd t ttjddgddgd t ttjddgddgd t ttjddgddgd d S )	Nr?   r   rY   g      @r  r  g      %@g      /@)r2   r3   r   rS  	TypeErrorr  r9   r9   r:   test_invalid_params'  s   z-TestMultivariateHypergeom.test_invalid_paramsN)r   r   r   r  r  r  r/   ZNINFr   rO  r!  r  r%  r   r:  re  rf  rl  rr  r/  r2  ry  rz  r;  r<  r=  r   r@  r9   r9   r9   r:   r  Z  sn    

"&&
	
r  c                 C   sd   | j }d| _ | j|ddi t| }| j|ddi}t|}|j|ddi}t|| || _ d S )Nr<   r   rP  )r   r   pickledumpsloadsr   )distfnrB  r\  rb   Zr0Z	unpickledZr1r9   r9   r:   check_pickling3  s   



rE  c                  C   s   t d} d| d< d| d< tdgtt dgfgtd| fgtd| fgtdg d	fgtd
gt	d
gg}|D ]\}}t
|| t|| q3d S )Nr+   r   r-   r  r9   r   r   r?   r  )r,   )r/   r1  r   r   r   r   r   r   r   r   r(   rE  )r   ZdistsrD  rB  r9   r9   r:   test_random_state_propertyG  s   


	
rF  )F__doc__rA  Znumpy.testingr   r   r   r   r   r   r  r   r2   Ztest_continuous_basicr
   r   r/   Zscipy.linalgre   Zscipy.stats._multivariater   r   r   r   Zscipy.statsr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   r%   Zscipy.integrater&   Zscipy.specialr'   Zcommon_testsr(   Zunittest.mockr)   r*   r   r   r,  r.  rF  r  r  r  r  r  r  r  rE  rF  r9   r9   r9   r:   <module>   sP     d   & 1 B @  %E][. o Z