o
    ?Hh"r                    @   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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-m.Z.m/Z/m0Z0m1Z1m2Z2 dd	lm3Z3m4Z4 dd
lm5Z5 ddl6m7Z7 ddl6m8Z8m9Z9m:Z:m;Z; ddl<m=Z= ddl>m?Z? ddl@mAZA ddlBmCZC dd ZDG dd dZEdBddZFdd ZGG dd dZHG dd dZIG dd dZJd d! ZKG d"d# d#ZLG d$d% d%ZMG d&d' d'ZNG d(d) d)ZOG d*d+ d+ZPG d,d- d-ZQG d.d/ d/ZRG d0d1 d1ZSG d2d3 d3ZTG d4d5 d5ZUG d6d7 d7ZVd8d9 ZWd:d; ZXG d<d= d=ZYG d>d? d?ZZG d@dA dAZ[dS )Cz8
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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random_tableuniform_directionvonmises_fisherdirichlet_multinomialvonmises)_covariance
Covariance)stats)tanhsinh)rombqmc_quaddblquadtplquad)multigammaln)check_random_state_property)_qsimvtv)patchc                 O   s@   t | t |} }t| |g|R i | t| j|j d S N)npasarrayr   r   shape)resrefargskwargs r>   c/home/air/sanwanet/gpt-api/venv/lib/python3.10/site-packages/scipy/stats/tests/test_multivariate.pyassert_close,   s   r@   c                   @   s  e Zd Zdd Zejejjejjejj	dd dZ
eee
Zeg dg dg dg d	geg d
g dg dg dgdZeedd eg d edd dZejdedd dd Zejdeeejdedd Zejdde ddgejdeeejdedd Zejde dgejdedd Zd d! Zejd"d#d$ Zd%d& ZdS )'TestCovariancec                 C   s  d}t jt|d ttd W d    n1 sw   Y  d}t jt|d tjtdtdd W d    n1 sBw   Y  d}t jt|d td W d    n1 s`w   Y  d	}t jt|d t	td W d    n1 sw   Y  d
}t jt|d t
dtdf W d    n1 sw   Y  d}t jt|d t
tddf W d    n1 sw   Y  d}t jt|d t
g dtdf W d    d S 1 sw   Y  d S )Nz:The input `precision` must be a square, two-dimensional...match   z0`precision.shape` must equal `covariance.shape`.   )
covariancez7The input `diagonal` must be a one-dimensional array...alpacaz9The input `cholesky` must be a square, two-dimensional...z4The input `eigenvalues` must be a one-dimensional...z,The input `eigenvectors` must be a square...z9The shapes of `eigenvalues` and `eigenvectors` must be...r	   rD   rE   )pytestr   
ValueErrorr*   CovViaPrecisionr7   oneseyeCovViaDiagonalCovViaCholeskyCovViaEigendecompositionselfmessager>   r>   r?   test_input_validation4   s8   "z$TestCovariance.test_input_validationc                 C   s   t | ddS )NTallow_singular)r   xr>   r>   r?   <lambda>V      zTestCovariance.<lambda>)Diagonal	PrecisionCholeskyEigendecompositionPSDrH   )   r	   rE   )r	         )rE   rb      )r	   r   rE   )r`   r   )rd   r`   r   r   r   r   )diagonal full rankzgeneral full rankzdiagonal singularzgeneral singularr	   N)r   rd   rg   cov_type_namerd   c           	      C   s   t g d}g d}ttd| }| j| }ttd|  }|||}|||}t|t|u s8J t|	||	| d S )NrH   )rD   r`   CovViafrom_)
r7   diaggetattrr*   _covariance_preprocessingr+   lowertyper   whiten)	rR   rh   ArX   cov_typepreprocessingfactoryr:   r;   r>   r>   r?   test_factoriesb   s   
zTestCovariance.test_factoriesmatrix_typec                 C   s  d| d| d}|| j | vrt| | j| }ttd| }| j| }t|dd}|||}t|j	|j	 t
|j|j t
|jt|j t|jt| tjd}	|	jdd}
||
}|
|j }t|| ||  t|d	rd
|vrt|||
 |	jdd}
||
}|
|j }t|d jdd|d jdd t|d	rd
|vrt|||
 t|d	r|tt|}t|j| | d S d S )Nrj    does not support 	 matricesTrU      VGiVK rE   size	_colorizesingularrD   rb   rE   rD   rd   axis)
_cov_typesrI   skip	_matricesrm   r*   rn   r   r@   log_pdetr   rankr9   r7   r8   rF   randomdefault_rngrq   UhasattrcolorizesumrM   lenT)rR   rw   rh   rS   rr   rs   rt   psd
cov_objectrngrX   r:   r;   r>   r>   r?   test_covariancep   s:   






"
zTestCovariance.test_covariancer|   r   c                 C   s  d| d| d}|| j | vrt| | j| }ttd| }| j| }g d}|||}	t}
t||dd}t||	dd}tj	
d}|j|||d}tj	
d}|
j||	||d	}tj	
d}|
||	|d
j|d}t|	tjrt|t| t|t| nt|j|j t|j|j t|| t|
|||	|| t|||| t|
|||	|| t|||| t|
||	|  t| |  d S )Nrj   rx   ry   皙?皙?333333?TrU   rz   r{   r|   random_stateseed)r   rI   r   r   rm   r*   rn   r   r7   r   r   rvs
isinstance	CovViaPSDr@   squeezer   r9   pdflogpdfentropy)rR   r|   rw   rh   rS   rr   rs   rt   meanr   mvndist0dist1r   rX   x1x2r>   r>   r?   test_mvn_with_covariance   s:   



z'TestCovariance.test_mvn_with_covariancec                 C   s   d}| j | }ttd| }| j| }g d}|||}t}	t||dd}
t||dd}tjd}|j|||d}t|		||||
	| t|	||
	| t|	
||||

| t|
||

| d S )Nrf   rj   r   TrU   rz   r{   )r   rm   r*   rn   r   r7   r   r   r@   cdflogcdf)rR   r|   rh   rw   rr   rs   rt   r   r   r   r   r   r   rX   r>   r>   r?   test_mvn_with_covariance_cdf   s   

z+TestCovariance.test_mvn_with_covariance_cdfc                 C   s>   d}t jt|d t  W d    d S 1 sw   Y  d S )Nz7The `Covariance` class cannot be instantiated directly.rB   )rI   r   NotImplementedErrorr+   rQ   r>   r>   r?   test_covariance_instantiation   s   "z,TestCovariance.test_covariance_instantiationzignore::RuntimeWarningc                 C   s   t g d}|jd }t |}tjtdd t||  W d    n1 s*w   Y  d}t j	
|}t j	
|}tt j|}t||}|j|d}	tj|||d}
t|	|
 d S )N)r	   rD   g:0yEr   zThe input matrix must be...rB   l   .ypGw r   )r7   rl   r9   zerosrI   r   rJ   r   r   r   r   r+   from_eigendecompositionlinalgeighr   )rR   rr   nr   r   rng1rng2covrvr:   r;   r>   r>   r?   test_gh9942   s   


zTestCovariance.test_gh9942c                 C   s   t d}tt dt df}tjj||d}|j	d d}t
|| tjjt ddgt ddgddggf}tjj||d}|j	d d}|d |d ksTJ |d |d ks^J d S )	NrD   r   r   r{         ?        g      y@r   r	   )r7   rL   r+   r   r   rM   scipyr,   r   r   r   array)rR   r   r   distr   r>   r>   r?   test_gh19197   s   

"zTestCovariance.test_gh19197)__name__
__module____qualname__rT   r7   rl   r   invcholeskyr   rn   r   list_all_covariance_typesr   r   rI   markparametrizerv   r   tupler   r   r   filterwarningsr   r   r>   r>   r>   r?   rA   2   sF    



/%

rA   Fc           	      C   s\   | | | f}||j }tj|\}}|r"|j| ddk}d||< |t| |j }|S )Nr{   r   )r   r   r7   r   r   normalrl   )	dimevalsr   r~   rr   _v	zero_eigsr   r>   r>   r?   _random_covariance  s   
r   c                 C   s$   t j| | }tj|\}}}|S r6   )r7   r   randnr   r   svd)r   Musr   r>   r>   r?   _sample_orthonormal_matrix   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ejdededdggd d! Zd"d# Zd$d% Zd&d' Zd(d) Zd*d+ Z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(ejdBdCdDdE Z)dFdG Z*ejdHe+dIe+dJgdKdL Z,ejdMe+dNe+dOe+dPgdQdR Z-dSdT Z.dUdV Z/dWdX Z0dYS )Z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 )NrE   rD   r   r	   r   r	   rD   )r7   arangeidentityassert_raisesrJ   r   r   r   )rR   mur   r>   r>   r?   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`   )
r7   r   r   r   r   r   ndimr   absr   )rR   rX   r   r   r   r   r>   r>   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`   )
r7   r   r   r   r   r   r   r   r   logrR   rX   r   r   d1d2r>   r>   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	   )	r7   r   r   r   r   r   r   r   r   rR   rX   r   r   d3d4r>   r>   r?   test_logpdf_default_valuesT     

z1TestMultivariateNormal.test_logpdf_default_valuesc                 C   r   r   )
r7   r   r   r   r   r   r   r   r   r   r   r>   r>   r?   test_logcdfa  r   z"TestMultivariateNormal.test_logcdfc                 C   r   r   )	r7   r   r   r   r   r   r   r   r   r   r>   r>   r?   test_logcdf_default_valuesk  r   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   rb   r	   TrU   )r7   r   r   r   rangedotr   r   r   r   r   )rR   r   r   expected_rankr   r   distnr>   r>   r?   	test_rankx  s   z TestMultivariateNormal.test_rankc              	   C   s  t ddD ]}tj|}t d|D ]}tj||}t||j}t||f}||d |d |f< t|}|d | |d |< t|}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|| |
|d d df  }||}||}t	|d t	|tj  qqd S )Nr	   r`   TrU   rd   r   )r   r7   r   r   r   r   r   r   r   r   r   r   r   r   r   inf)rR   r   zkr   cov_kkcov_nnrX   r   cov_rrydistn_kkdistn_nndistn_rrpdf_kkpdf_nnpdf_rr	logpdf_kk	logpdf_nn	logpdf_rry_orthpdf_rr_orthlogpdf_rr_orthr>   r>   r?   test_degenerate_distributions  sR   











z4TestMultivariateNormal.test_degenerate_distributionsc           
      C   s   d}t ddD ]\}t d|D ]T}t|}t|d d d |f }t||j}tj|||d}tj|||dd}t	|j
| t|dksHJ tj|||dd}	t	|	j
| t|	tj ksbJ qqd S )	N
   rD   ra   r	   r   r   r|   T)r   r   rV   r   )r   r7   r   r   r   r   r   r   r   r   r|   allr   r   )
rR   r   r   rmnr   vrXr   r   r>   r>   r?   test_degenerate_array  s&   


z,TestMultivariateNormal.test_degenerate_arrayc                 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   rD   dtyper   r	   )r7   expr   floatfill_diagonalr   r   r   detr   r   slogdetr   r   )rR   large_total_lognposnzerolarge_entryr   r   r   r>   r>   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   rb   rD   rE   MbP?rtol)r7   r   r   r   r   r   r   r   r   r   r   )rR   r   datar   r   r  desired_pdfdesired_cdfijactualr>   r>   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   rD   r  )333333??      ?)r7   linspacer   r   r   r   r   )rR   rX   r   r   scaler   r   r>   r>   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      @r1  r   333333?i  ra   r	   r   rD   r   r   r   locr3  )r	   r	   {Gz?r'  atol)
r7   r   r2  meshgridemptyr   r   r.   r   r   )rR   r   r   r   deltar   xvyvposr   margin_xmargin_ygauss_xgauss_yr>   r>   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 r   )r7   r   r   r   r   r   r   r   r   r   r   )rR   rX   r   r   norm_frozenr>   r>   r?   test_frozen$  s   


z"TestMultivariateNormal.test_frozenrF   rD   r	   c                 C   sF   t d}t d}t||}t |j|sJ t |j|s!J d S )NrD   rD   )r7   rL   rM   r   allcloser   r   )rR   rF   r   cov_should_berG  r>   r>   r?   2test_frozen_multivariate_normal_exposes_attributes2  s
   


zITestMultivariateNormal.test_frozen_multivariate_normal_exposes_attributesc           	   	   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   rc   r1  r   r   Hz>rd   h㈵>)cond)r7   r   r   r   r   r   r   r   r   fullrl   r   pinvr   r   r   r   )	rR   r   rX   r   r   r   rO  r   psd_pinvr>   r>   r?   test_pseudodet_pinv@  s    z*TestMultivariateNormal.test_pseudodet_pinvc                 C   s    g dg dg}t tt| d S )NrH   rb   r`   ra   r   rJ   r   rR   r   r>   r>   r?   test_exception_nonsquare_covZ     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   )r7   nanr   rJ   r   r   )rR   cov_nancov_infr>   r>   r?   test_exception_nonfinite_cov^  s   z3TestMultivariateNormal.test_exception_nonfinite_covc                 C   s    ddgddgg}t tt| d S )Nr	   r   rd   rU  rV  r>   r>   r?   test_exception_non_psd_covd  rX  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dgddgg}d}tjt jj|d t|d W d    d S 1 sjw   Y  d S )	Nr   r`   )r`   r`   r   r   z0When `allow_singular is False`, the input matrixrB   r   )r7   r   r   r   rL   r   LinAlgErrorr   r   r   r   r   r   rI   r   )rR   rX   r   r   emsgr>   r>   r?   test_exception_singular_covh  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   rD   r`   rE   )r	   rE   rD   dr	   rD   r   )rD   r`   r1  )r   r1  rE   绽|=r;  )gwk:E]?g̙Z?g0?gL8?gAG*?gh㈵>)gIى?gG?gƏ)?gėh?gEE?r	   rN  )	r7   r   r2  cosr   r   r   r   r   )rR   r_pdfrX   r  r   r  r   r   r   r_cdfr   r_cdf2r2mean2cov2cdf2r>   r>   r?   test_R_values{  s$   z$TestMultivariateNormal.test_R_valuesc                 C   s<   t d}t d}t||dd}| }t|ddg d S )NrD   rD   rD   TrU   r   )r7   r   r   r   r   )rR   r   rF   modelsampler>   r>   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,  rb   r	   r  rD   r   r   r   )r   r   r7   r   r   r9   r   )rR   Nrc  rr  r   r>   r>   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  rE     r   r&  r   )r7   r   r   r   r   r   r   r   r   r   r   )rR   r   r   r   r   r|   rr  r>   r>   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 )Nrv  rE   r   r1  rD   r	   )r7   r   r   r   r   r   r   r   r   r   eigr   pir   )rR   r   r   r   r   r   eigsdesiredr>   r>   r?   test_entropy  s   
,z#TestMultivariateNormal.test_entropyc                 C   s*   t g d}d}tt t|| d S )Nr	   r	   r	   r1  )r7   r   r   r  r   )rR   alphar|  r>   r>   r?   test_lnB  s   zTestMultivariateNormal.test_lnBc                 C   s   t jd}ddg}t d}|dd d }|dd d }tj||||d}t|||}t|||}t j|ddd	f |dd	df fd
d}	t j|dd	df |ddd	f fd
d}
t|	||}t|
||}|| | | }t|| d S )N   FYc@Y r   rD   )rb   rE   rD   ra   rE   lower_limit.r	   rd   r   )r7   r   r   rM   r   r   concatenater   )rR   r   r   r   abcdf1cdf2acdf2bab1ab2cdf2ab1cdf2ab2rn  r>   r>   r?    test_cdf_with_lower_limit_arrays  s   
**z7TestMultivariateNormal.test_cdf_with_lower_limit_arraysc           
      C   s   t jd}|d}|d}||j }|dd d }|dd d }tj||||d}t||j||d}t tj||||d}t t||j||d}	t||dd t||dd t|	|dd d S )	Nr  rE   rE   rE   rD   rE   ra   r  -C6?r&  )	r7   r   r   r   r   r   r  r   r   )
rR   r   r   r   r  r  r  rn  cdf3cdf4r>   r>   r?   %test_cdf_with_lower_limit_consistency  s   


z<TestMultivariateNormal.test_cdf_with_lower_limit_consistencyc                 C   s   t d}t d}g dg dg dg dg}g dg dg dg dg}t g d}tj||||d}t||d |  d S )	NrE   r~  re   r	   r   r	   r   r	   r   r	   rd   rd   r	   r  r   r7   r   rM   r   r   r   r   )rR   r   r   r  r  expected_signsr   r>   r>   r?   test_cdf_signs"  s   

z%TestMultivariateNormal.test_cdf_signsc                 C   s   t dt g d }t|}d}tjt|d t	ddg| W d    n1 s-w   Y  tjt|d tddg| W d    n1 sKw   Y  g d}t
|g d|}ttj
||d| t
|g d	|}ttj
|d|d| d S )
Nr	   rH   z7`cov` represents a covariance matrix in 3 dimensions...rB   r   )r1  r1  r1  re   r^  r~  )r7   rl   r   r*   rK   rI   r   rJ   r   r   r   r   )rR   Pr   rS   rX   r;   r>   r>   r?   test_mean_cov-  s   
z$TestMultivariateNormal.test_mean_covc                 C   sJ   ddg}d}t jt|d t| W d    d S 1 sw   Y  d S )Nr	   rE   z`x` must be two-dimensional.rB   )rI   r   rJ   r   fit)rR   r(  	error_msgr>   r>   r?   test_fit_wrong_fit_data_shape@  s
   "z4TestMultivariateNormal.test_fit_wrong_fit_data_shaper   )rE   r`   c                 C   sh   t jd}|d|f}t|\}}t j|ddt j|jdd}}t||dd t||dd d S )	N   @,Q| r  r   r   )ddofV瞯<rf  r&  )	r7   r   r   r   r  r   r   r   r   )rR   r   r   rX   mean_estcov_estmean_refcov_refr>   r>   r?   test_fit_correctnessF  s    z+TestMultivariateNormal.test_fit_correctnessc                 C   sF   t dd}d}t d}tj|||d\}}t|| t|| d S )N)rD   r	   rE   r   )fix_meanfix_cov)r7   rP  
atleast_2dr   r  r   )rR   r(  
mean_fixed	cov_fixedr   r   r>   r>   r?   test_fit_both_parameters_fixedO  s   


z5TestMultivariateNormal.test_fit_both_parameters_fixedr  rp  rE   c                 C   L   d}t jt|d tjtd|d W d    d S 1 sw   Y  d S )Nzd`fix_mean` must be a one-dimensional array the same length as the dimensionality of the vectors `x`.rB   rD   r  rI   r   rJ   r   r  r7   rM   )rR   r  ra  r>   r>   r?   "test_fit_fix_mean_input_validationX  s   "z9TestMultivariateNormal.test_fit_fix_mean_input_validationr  rI  )rE   rD   )rb   rb   c                 C   r  )Nzn`fix_cov` must be a two-dimensional square array of same side length as the dimensionality of the vectors `x`.rB   rE   r  r  )rR   r  ra  r>   r>   r?   +test_fit_fix_cov_input_validation_dimension`  s   "zBTestMultivariateNormal.test_fit_fix_cov_input_validation_dimensionc                 C   sb   d}t jt|d tddgddgg}tjtd|d W d    d S 1 s*w   Y  d S )Nz2`fix_cov` must be symmetric positive semidefinite.rB   r   r   g      rD   r  )rI   r   rJ   r7   r   r   r  rM   )rR   r  r  r>   r>   r?   *test_fit_fix_cov_not_positive_semidefinitej  s
   "zATestMultivariateNormal.test_fit_fix_cov_not_positive_semidefinitec                 C   s   t jd}|d}|d}t ||j}tj||d|d}t|\}}tj|||d	 }tj||d\}	}
t
|	| tj||	|
d	 }||k sOJ |d}dt ||j }|
| }tj||	|d	 }||k sqJ d S )	Nr  rE   r  r  r   r   r|   r   r   r  :0yE>)r7   r   r   r   r   r   r   r  r   r   r   )rR   r   r8  rr   r   samples	mean_freecov_free	logp_freemean_fixcov_fixlogp_fixmcov_perturbedlogp_perturbedr>   r>   r?   test_fit_fix_meanp  s>   




z(TestMultivariateNormal.test_fit_fix_meanc                 C   s   t jd}|d}|d}t ||j}tj||d|d}t|\}}tj|||d	 }tj||d\}	}
t
|	t j|dd	 t
|
| tj||	|
d	 }||k sYJ |	d
|d  }tj|||
d	 }||k srJ d S )Nr  rE   r  r  r  r   r  r   r   r  )r7   r   r   r   r   r   r   r  r   r   r   r   )rR   r   r8  rr   r   r  r  r  r  r  r  r  mean_perturbedr  r>   r>   r?   test_fit_fix_cov  s<   


z'TestMultivariateNormal.test_fit_fix_covN)1r   r   r   r   r   r   r   r   r   r   r  r  r$  r.  r4  rF  rH  rI   r   r   r7   rM   r+   from_diagonalrL  rS  rW  r\  r]  rb  ro  rs  ru  rx  r}  r  r  r  r  r  r  r  r  r   r  r  r  r  r  r>   r>   r>   r?   r   &  sj    

4
@
	

r   c                   @   sL   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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j
||t ||f t|tj
|t ||f| t|t||t ||f t|t|t ||f| d S )Nrb   rE   r   r1  ffffff?)r`   rb   rE   r  )r7   rP  r   r   rJ   r   r   r   r   r_  r   rL   )rR   num_rowsnum_colsr   r   Vr`  r>   r>   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 )Nrb   rE   r   r1  r  r	   r   rowcovcolcov)r   )r  )r  )r   r  )r   r  )r  r  )r7   rP  r   r   r   r   r   r9   r  r  r   )rR   r  r  r   r   r  ZZrZcIrIcI1r>   r>   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 )Nrb   rE   r   r   r   r  )r7   rP  r   r   r   r  r  )
rR   r  r  r   UvUsVvVsr  r  r>   r>   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   r1  r  r  r   r   )r   r  r  r   )	r   r7   rP  r   r   r   r   r   r   )rR   r+  r,  r   r   r  frozenrvs1rvs2r  pdf1pdf2logpdf1logpdf2r>   r>   r?   test_frozen_matrix_normal  s(   




z*TestMatrixNormal.test_frozen_matrix_normalc                 C   s  t ddD ]}t ddD ]}}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j||d	}t||d
d t|	|d
d t|
| qqd S )Nr	   r`   r   r1  r  r  r   r   r   re  r&  )r   r7   rP  r   r   r   r   r   r   r   flattenkronr   r   )rR   r+  r,  r   r   r  r  r  r  r  entropy1vecXvecMr   r  r  entropy2r>   r>   r?   test_matches_multivariate  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 )Nrb   rE   r   r1  r  r  r  r   r   i  r   r   rD   re  )r7   rP  r   r   r   r  newaxisr   r9   r   r   r   )rR   r  r  r   r   r  rt  r  X1X2r  array_logpdfr+  r,  separate_logpdfr>   r>   r?   test_array_input4  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 )Nrb   rE   r   r1  r    r  r   r   r   r   r   rf  r	   rD   )r7   rP  r   r   r   r   r   r   reshaper   swapaxes)rR   r  r  r   r   r  rt  r  r  sample_meansample_colcovsample_rowcovr>   r>   r?   test_momentsK  s$   zTestMatrixNormal.test_momentsc                 C   s   t jtddgddggtddgddggtddgddggtjddd	}td
dgddggddgddggg}t|| d S )Nr	   rD   rE   rb   rd   r`   r  r   )r   r  r  r   r|   g6?gǀgcϿ@g"'@gI2͖@g^jy%@gDn3@g|\3@)r   r   r7   r   r   r   r   )rR   r-  expectedr>   r>   r?   test_samplesa  s"   
zTestMatrixNormal.test_samplesN)r   r   r   r  r  r  r  r  r  r  r  r>   r>   r>   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 Nrv  r	       &.>r  r  rd   )r7   r   r   randintr   r   r   varr   r   r   r   r   r   rR   r   r  rc  	num_testsr+  rX   r>   r>   r?   test_frozen_dirichletv  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 )Nrv  r          @      @rc   r{   )rc   rE   rd   )r7   r   r   r   r   r   r9   r   rJ   r   r   r   rR   r  rX   r>   r>   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 )Nrv  )r   r   r  r  rc   r{   
r7   r   r   r   maximumr   r   rJ   r   r   r  r>   r>   r?   test_alpha_with_zeros  
   z#TestDirichlet.test_alpha_with_zerosc                 C   r  )Nrv  )r   g       r  r  rc   r{   r  r  r>   r>   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   ra   )r7   r   r   r   r   r   r   r  r>   r>   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   r1  r  r  r  r7   r   r   rJ   r   r   r   r  r>   r>   r?   $test_data_with_zeros_and_small_alpha     z2TestDirichlet.test_data_with_zeros_and_small_alphac                 C   r  )Nr  )r   皙r   r  r  r  r>   r>   r?   test_data_with_negative_entries  r  z-TestDirichlet.test_data_with_negative_entriesc                 C   r  )Nr  )r   皙?r   r  r  r  r>   r>   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  )rD   rc   rc   g$I$I?r7   r   rP  r   rJ   r   r   r   r  r>   r>   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  )rD   rD   rc         ?r  r  r>   r>   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  rE   rc   UUUUUU?)r7   r   rP  r   r   r   r  r>   r>   r?   test_alpha_correct_depth  s   z&TestDirichlet.test_alpha_correct_depthc                 C   r  )Nr  r  r1  r  r  r>   r>   r?   test_non_simplex_data  r  z#TestDirichlet.test_non_simplex_datac                 C   r  )Nr  )rD   rc   r1  r  r  r>   r>   r?   test_data_vector_too_short  r  z(TestDirichlet.test_data_vector_too_shortc                 C   r  )Nr  )r`   rc   r   r  r  r>   r>   r?   test_data_vector_too_long  r  z'TestDirichlet.test_data_vector_too_longc                 C   sj   t g d}t|}g d}g d}g dg dg dg}t| | t| | t| | d S )N)r   皙?r   r1  皙?r   )UUUUUU?{Gz?Q?)r(  )r+  r)  O贁N)r,  r-  r*  )r7   r   r   r   r   r  r   )rR   r  rc  expected_meanexpected_varexpected_covr>   r>   r?   test_mean_var_cov  s   zTestDirichlet.test_mean_var_covc                 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   )	r7   r   r   r   r   r   r  r   r   )rR   r  rc  r>   r>   r?   r     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  )
r7   r   r   r  r   r   r   r   r   r   r  r>   r>   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 )Nrv  r	   r  r  r  r  r`   )r7   r   r   r  r   r   r   r   vstackr   r   appendr   )rR   r   r  rc  r   num_multiplexmr+  r  rX   rmrsxsr  r>   r>   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 )Nrv  r  r  rD   r   r	   r  )r7   r   r   r   r   r   r   r   r   r   r   r  )rR   r  rc  r  r   r+  rX   r>   r>   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$  r1  r   r2  r:  r;  r>   r>   r>   r?   r  t  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)r7   r   r   rJ   r   r   strr   )r   sigmar`  ra  r>   r>   r?   ,test_multivariate_normal_dimensions_mismatch=  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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	   rD   ndminr   r	   rD   r  rE   )
r7   r   r_r   r   r3  r9   r   rJ   rM   )rR   
true_scalescalesr3  wr>   r>   r?   test_scale_dimensionsQ  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	   rD   r@  rE   rH   rp  r   )r7   r   rC  r   r   r   r  rM   )rR   r  rF  densityrX   r>   r>   r?   test_quantile_dimensions|  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 	Nrb   r	   rD   rd   r   r`   r  r   )r7   rl   r   tril_indicesr   r   r   r4  r   r2  r   r   r  r   moder   r   )	rR   r   r3  r  r+  rX   
parametersdfrF  r>   r>   r?   rH    s,   $$
zTestWishart.test_frozenc                 C   s   d}d}t |}d|d< d|d< t||}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|| t|| d S )NrE   r  r1  r   r	   r   i r{   r	   rD   rd   rK  )r7   rM   r   r   r   r   r   rC  	chisquarerl   rL  r   r   r   r   r   )rR   r   rO  r3  rF  w_rvsfrozen_w_rvscovariances	variancesrr   DDAmanual_w_rvsr>   r>   r?   test_wishart_2D_rvs  s6   





zTestWishart.test_wishart_2D_rvsc                 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  rD   r  r   numr{   r9  r   )r7   r   r   rM   r   r  r2  r   r   r   r  r   r   r   r   r
   )rR   snr   r3  df_ranger  rO  rF  cr   r<   r  r>   r>   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 )NrZ  r[  r  rb   r	   ra   rd   rK  r3  r   r\  r{   r   r9  r   )r7   r   r   rl   r   rL  r   r   rL   r   r   r   r   r  r   r   r2  r   r   r
   )rR   r^  rO  r   r3  lamdasigma_lamdarF  r`  r  r   r<   r  r>   r>   r?   test_is_scaled_chisquared  s(   

z%TestWishart.test_is_scaled_chisquaredN)	r   r   r   rG  rI  rH  rY  ra  re  r>   r>   r>   r?   r?  P  s    +1 ,r?  c                   @   s   e Zd Zdd Zdd Zdd Zejddd	gd
d Z	dd Z
dd Zejdddgdd Zdd Zejdddgdd Zdd Zejdddgdd Zdd Zdd Zd d! Zd"S )#TestMultinomialc                 C   s   t ddd}t|ddd t ddgd	d
dg}|tj ks!J t d	d	gd	d
dg}|d	ks2J t ddgd	ddg}t|tjdd d S )NrE   rb   rc   r   r  g&Mtyr  r&  rE   rb   r   r   r  rg   )r   logpmfr   r7   r   rY  )rR   vals1vals2vals3vals4r>   r>   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  rc   rh  rE   r   r  r&  )ra         r   r0  ra   r   )r   ri  r    r   pmfrR   val1val2r>   r>   r?   test_reduces_binomialI  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 )NrE   )g      ?r  g      ?g     @?g     ?g      ?g      `?g     ?g      ?g      ?g      ?g      ?)
)r   r   rE   )r	   r   rD   )rD   r   r	   )rE   r   r   r   r~  rD   r	   r   )r   rD   r	   rd  )r   rE   r   +=rf  )r   r   rr  )rR   r   pr_valsrX   r>   r>   r?   test_RT  s   zTestMultinomial.test_Rr   r   rE   c                 C   sH   t j|dgd ddd}tjd}|j |dgd dd}t|| d S )Nr  rE   rc   {   r   r{   )r   r   r7   r   RandomStater   )rR   r   sc_rvsrndmnp_rvsr>   r>   r?   test_rvs_npe  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 t g ddg d}|dksJ t g ddg d}|dksJ d S )N)r`   r`   r	   r	   r  r&  rg  rc   rh  g5
?rE   r   ro  rd   	   rq  g*7?g<Wƌ?r   rD   r  rb   rD   rE   rE   r   ra   )gUUUUUU?r   r   gߦ?re   rw  )r   rr  r   r7   r=  float64r   )	rR   vals0rj  rk  rX   rl  rm  vals5vals6r>   r>   r?   test_pmfm  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	   rD   rE   r   r0  r   r%  gv?g~jt?r  r&  rb   r   )r   rr  r   )rR   r  rj  rk  rl  rm  r>   r>   r?   test_pmf_broadcasting  s     &z%TestMultinomial.test_pmf_broadcastingr`   c                 C   s   t |d}|d d | d d | d d g| d d |d d | d d g| d d | d d |d d gg}t||dd d S )	N)r   r   r1  r   r%  r   r1  r  r  r&  r   r   r   )rR   r   cov1rm  r>   r>   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   r0  r   r%  g?gܿgr  r&  rb   g
ףp=
?g
ףp=
׿r   r  r'  r5  gzG?gzGr/  g333333r  rR   r  rm  cov3cov4cov5cov6r>   r>   r?   test_cov_broadcasting  s     z%TestMultinomial.test_cov_broadcastingrD   c                 C   s*   t |ddg}t|t|ddd d S Nr   r%  r  r&  r   r   r   r    )rR   r   ent0r>   r>   r?   r}    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 )NrD   rE   r   r   r  r&  rc   ro  r  r'  r5  r  )rR   r  ent1ent2r>   r>   r?   test_entropy_broadcasting  s    
z)TestMultinomial.test_entropy_broadcastingc                 C   s.   t |ddg}t||d |d gdd d S r  r   r   r   )rR   r   mean1r>   r>   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`   ra   r   r%  r   r  g433333?g433333@r  r&  r  )rR   r  r>   r>   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	   rD   ro  r  )
r7   r   r   r8   r  r   r   rr  ri  r   )rR   r   pvalsrX   	mn_frozenr>   r>   r?   rH    s   "
zTestMultinomial.test_frozenc                 C   sZ   d}t jd}||}d|d< |t | }t |}t|||}t |s+J d S )NX   l   iRZY gKH9rd   )r7   r   r   r   rL   r   ri  isfinite)rR   r   r   ry  rX   ri  r>   r>   r?   test_gh_11860  s   

zTestMultinomial.test_gh_11860N)r   r   r   rn  rv  r{  rI   r   r   r  r  r  r  r  r}  r  r  r  rH  r  r>   r>   r>   r?   rf  ;  s&    



rf  c                   @   4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )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 rJ  )r7   rl   r   rL  r   r   r   r4  r   r2  r   r   r  r   rM  r   r   )	rR   r   r3  r  r+  rX   rN  rO  iwr>   r>   r?   rH    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 ]K}t||}t|d dd}t	|
 |
  t	| |  t	|||| |j|d}	|d ddf}
d}td|
||	 t	| |  q"d S )NrZ  r[  r	   r`      rD   r  r   r  r\  r1  rb  r{   r   r9  r   )r7   r   r   rM   r   r  r2  r   r   r   r  r   r   r   r
   r   )rR   r^  r   r3  r_  r  rO  r  igr   r<   r  r>   r>   r?   test_1D_is_invgamma  s$   

z"TestInvwishart.test_1D_is_invgammac                 C   s  d}d}t |}d|d< d|d< t||}t jd t||}t jd | }t jd t jjdd}t jt j|d t j|d	 t j|f d }t 	|}	||	t j
|d
d< t j|}
t j|	j|
jj}t ||j}t|| t|| d S )NrE   r  r1  r   rP  iHG	 r{   rD   r	   rd   rK  )r7   rM   r   r   r   r   r   rC  rQ  rl   rL  r   r   solver   r   r   )rR   r   rO  r3  r  iw_rvsfrozen_iw_rvsrT  rU  rr   rV  Lmanual_iw_rvsr>   r>   r?   test_invwishart_2D_rvs)  s6   

	


z%TestInvwishart.test_invwishart_2D_rvsc                 C   s   d}d}dD ]h}t t |d }t ||d  d |t j|dd< t |j|}t||}| }| }|| d }|j	|d	d
}	|	jdd}
||d  d }d| }t
d|d  }t j|
| | d|dsnJ qdS )z1Test that sample mean consistent with known mean.r   N  )r	   r`   r	   rD   rd   rK  r1  r   r   r   r   r9  rf  N)r7   rl   r   rL  r   r   r   r   r  r   r   ppfrJ  )rR   rO  sample_sizer   r3  r   	Xmean_expXvar_exp	Xmean_stdr  	Xmean_estntests	fail_ratemax_diffr>   r>   r?   test_sample_meanW  s*   $

zTestInvwishart.test_sample_meanc                 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 )zRegression test for gh-8844.)rD   r	   r   r1  )r	   rD   r1  r1  )r   r1  rE   r	   )r1  r1  r	   rD   )r  rc   rE   r	   )rc   r  r`   r	   )rE   r`   ro  rD   )r	   r	   rD   r  ra   r   rD   r	   r1  N)r7   r   r   r   r9   r   r  r  r   r2   tracer   )rR   r  Psinuprobry  siglogdetX	logdetPsir   r  r>   r>   r?   test_logpdf_4x4r  s4   




zTestInvwishart.test_logpdf_4x4N)r   r   r   rH  r  r  r  r  r>   r>   r>   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  rE   )g(fg^A.g*B?)g:d?g2ugi}?)gR#xlF?g^`D?g?r   r   )r7   r   r   r   r   r   r   r}  )rR   rX   r  r   r>   r>   r?   test_reproducibility  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 Nrp  r	   r   )r   rJ   r   r   rR   r>   r>   r?   test_invalid_dim     z&TestSpecialOrthoGroup.test_invalid_dimc                 C   s4   d}t |}|jdd}t j|dd}t|| d S )Nrc   r   r   )r   r   r   )rR   r   r  r  r  r>   r>   r?   test_frozen_matrix  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   .0r   r+  r>   r>   r?   
<listcomp>      z<TestSpecialOrthoGroup.test_det_and_ortho.<locals>.<listcomp>rD   r  c                 S      g | ]}t j|qS r>   r7   r   r  r  rX   r>   r>   r?   r        r      vIh%<=r&  r   )r   r   r   r7   r   r   rM   r9   )rR   r9  detsrX   r>   r>   r?   test_det_and_ortho  s   z(TestSpecialOrthoGroup.test_det_and_orthoc                    x   d}d}d}t jd tj||dd fdd D  fd	d
 D }fdd
|D }t|gt| | d S )Nr`   r  皙?r  r{   r6  r  )r	   rb   r  c                    .   i | ]\  ft  fd dD qS )c                       g | ]}|   qS r>   r>   r  ecerr>   r?   r    r  z>TestSpecialOrthoGroup.test_haar.<locals>.<dictcomp>.<listcomp>sortedr  r9  r  r?   
<dictcomp>     . z3TestSpecialOrthoGroup.test_haar.<locals>.<dictcomp>c                    &   g | ]} D ]
}||kr||fqqS r>   r>   r  e0e1elsr>   r?   r       & z3TestSpecialOrthoGroup.test_haar.<locals>.<listcomp>c                    &   g | ]\}}t  |  | d  qS r  r   r  p0p1projr>   r?   r    r  )r7   r   r   r   r   r   r   rR   r   r  ks_probpairsks_testsr>   r  r  r9  r?   	test_haar     
zTestSpecialOrthoGroup.test_haarN)r   r   r   r  r  r  r  r  r>   r>   r>   r?   r    s    	r  c                   @   s^   e Zd Zdd Zdd Zdd Zdd Zej	d	g d
dd Z
dd Zejjdd ZdS )TestOrthoGroupc                 C   sn   d}t j| td}tjd|d}tt j|d t g dg dg dg}t	|| t	|| d S )Nr  rE   r   rd   )gm?gӠh"gEo?)g~7ݲC?g>Ŀg0J_ٿ)gݱ&ǿg 6qrg`X|[)
r7   r   r   r   r   r   r   r  r   r   )rR   r   rX   r   r  r>   r>   r?   r    s   


z#TestOrthoGroup.test_reproducibilityc                 C   r  r  )r   rJ   r   r   r  r>   r>   r?   r    r  zTestOrthoGroup.test_invalid_dimc                 C   V   d}t |}t |dd}|jdd}t j|dd}|jdd}t|| t|| d S )Nrc   r   r   r   r	   r{   )r   r   r   rR   r   r  frozen_seedr  r  rvs3r>   r>   r?   r       
z!TestOrthoGroup.test_frozen_matrixc                 C   s|   dd t ddD }tdd |D }tt|t|jdd |D ]}|D ]}tt||j	t
|jd  q(q$d S )	Nc                    s"   g | ]  fd dt dD qS )c                    s   g | ]}t  qS r>   )r   r   r  r+  r   r>   r?   r    s    @TestOrthoGroup.test_det_and_ortho.<locals>.<listcomp>.<listcomp>r  )r   r  r>   r
  r?   r    s
    

z5TestOrthoGroup.test_det_and_ortho.<locals>.<listcomp>rD   r  c                 S   s   g | ]	}d d |D qS )c                 S   r  r>   r  r  r>   r>   r?   r    r  r  r>   )r  xxr>   r>   r?   r        r  r&  r   )r   r7   r   r   fabsrL   r9   r   r   r   rM   )rR   r9  r  r  rX   r>   r>   r?   r    s   z!TestOrthoGroup.test_det_and_orthor   )rD   r`   r  r  c                 C   s   t jd}t|d}|jd|d}tj|}t |dk}t	|}t
||}|jdd\}	}
|	d  k r<|
k s?J  J d S )	Nl   <`Lr r
  rw  r   r   gffffff?)confidence_levelr1  )r7   r   r   r   r   r   r   r  r   r   r,   	binomtestproportion_ci)rR   r   r   r   r   r  r   r   r:   lowhighr>   r>   r?   test_det_distribution_gh18272   s   
 z,TestOrthoGroup.test_det_distribution_gh18272c                    r  )Nr`   r  r  i  r{   r  c                    r  )c                    r  r>   r>   r  r  r>   r?   r  #  r  z7TestOrthoGroup.test_haar.<locals>.<dictcomp>.<listcomp>r  r  r  r  r?   r  #  r  z,TestOrthoGroup.test_haar.<locals>.<dictcomp>c                    r  r>   r>   r  r  r>   r?   r  $  r  z,TestOrthoGroup.test_haar.<locals>.<listcomp>c                    r  r  r  r  r  r>   r?   r  %  r  )r7   r   r   r   r   r   r   r  r>   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{   )r7   r   r   r   r   r   )r   r   _sr   r>   r>   r?   random_ortho-  s    z<TestOrthoGroup.test_pairwise_distances.<locals>.random_orthorD   ra   r  re  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
  rD   )r7   r   )r  r   )r   r   r>   r?   r  3  s    z\TestOrthoGroup.test_pairwise_distances.<locals>.generate_test_statistics.<locals>.<listcomp>r{   )r7   r   r   r   r   r9   )r   rt  epsr,   r
  )r   r?   generate_test_statistics2  s
   zHTestOrthoGroup.test_pairwise_distances.<locals>.generate_test_statisticsr  )r  re  )
r7   r   r   r   r   r,   r   r   r   r   )rR   r  r  r  r-  _Dry  r>   r
  r?   test_pairwise_distances(  s   	z&TestOrthoGroup.test_pairwise_distancesN)r   r   r   r  r  r  r  rI   r   r   r  r  slowr  r>   r>   r>   r?   r    s    
r  c                   @   r  )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  r1  r%  r/  r   r   )r   ѓ2ǿA?Ϳ)r  r   4($?R}%?)r   r"  r   qTnƿ)r!  r#  r$  r   )r7   r   r   r   r   r   r   )rR   r{  rX   r   r  r>   r>   r?   r  D  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	   rD   rE   rb   g      )r	   rD   r   )r   rJ   r   r   r  r>   r>   r?   test_invalid_eigsP  s   z'TestRandomCorrelation.test_invalid_eigsc                 C   sR   d}t |}t |dd}t j|dd}|jdd}| }t|| t|| d S )Nr  r  r   r   )r   r   r   )rR   r{  r  r  r  r  r  r>   r>   r?   r  Y  s   
z(TestRandomCorrelation.test_frozen_matrixc           
         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 r6   )r   )r+  r`  r>   r>   r?   r   m  s   z3TestRandomCorrelation.test_definition.<locals>.normr|  c                    s    g | ]} |t jj|d qS )r{   )r7   r   r   r	  r   r>   r?   r  r  s     z9TestRandomCorrelation.test_definition.<locals>.<listcomp>rD   ra   )rb   r   r   r   c                 S   s   g | ]	}d gt | qS )r   )r   r  r`  r>   r>   r?   r  u  r  c                 S      g | ]}t |qS r>   )r   r   r(  r>   r>   r?   r  v      c                 S   s   g | ]}t t j|qS r>   )r7   r  r   r  r  r>   r>   r?   r  |  s    c                 S   r)  r>   )r7   prodr(  r>   r>   r?   r  }  r*  r  r:  c                 S   r)  r>   )r7   rl   r  r>   r>   r?   r    r*  r&  )r7   r   r   r   r4  r   zipr   )
rR   r{  rL   r9  r  
dets_knowndiagsr  r  rX   r>   r'  r?   test_definitione  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	   r  ignore)overrd   gu <7~r0  r  rD   r6  gMk   @)r7   r   r  r   _to_corrr   errstate	nextaftercopyr   r   )rR   r  gm0r>   r>   r?   test_to_corr  s$   
"z"TestRandomCorrelation.test_to_corrN)r   r   r   r  r&  r  r/  r8  r>   r>   r>   r?   r  C  s    	$r  c                   @   st   e Zd Zejdddgejdg ddd Zejdg dd	d
 Zdd Zejdg ddd Z	dS )TestUniformDirectionr   r	   rE   r|   Nr	   r`   )r`   rb   c           
      C   sv   t jd}t||d}||}t |t |}}|j|||dj}|j|ks,J t j	j
|dd}	t|	d d S )N   Sid)_i4 r   r{   rd   r   r   )r7   r   r   r&   r   r   rM   r   r9   r   r   r   )
rR   r   r|   r   uniform_direction_distr  r   r   expected_shapenormsr>   r>   r?   r    s   
z!TestUniformDirection.test_samples)Nr   rp  r   c                 C   sB   d}t jt|d t| W d    d S 1 sw   Y  d S )NzMDimension of vector must be specified, and must be an integer greater than 0.rB   )rI   r   rJ   r&   r   )rR   r   rS   r>   r>   r?   r    s   "z%TestUniformDirection.test_invalid_dimc                 C   sR   d}t |}t |dd}|jdd}t j|dd}| }t|| t|| d S )Nr`   r  r   r   )r&   r   r   r  r>   r>   r?   test_frozen_distribution  s   
z-TestUniformDirection.test_frozen_distribution)rD   r`   ro  c                 C   s   t jd}t||d}|jdd\}}||| | 8 }|t j| }t|| ddd |jdd}|| }|| }t ||}	|	t j	7 }	|	dt j	  }	t
 }
t|	|
j}|jd	ks]J d S )
Nl   @9Y sr   rD   r{   r   rx  rf  '  r  )r7   r   r   r&   r   r   r   r   arctan2rz  r   r   r   pvalue)rR   r   r   spherical_distv1v2r  s1s2anglesuniform_distkstest_resultr>   r>   r?   test_uniform  s   
z!TestUniformDirection.test_uniformN)
r   r   r   rI   r   r   r  r  r?  rK  r>   r>   r>   r?   r9    s    
r9  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  rE   r   )y~p?yZ?y뉦?ĭ?y8Z?&4?)yEHξr?8i?ye??+ݓݿyu=ٿB۽'?)yqXQÿ?yN6ҿzܷy! 8l?2W'?)r7   r   r   r   r   r   r   )rR   rX   r   r  r>   r>   r?   r    s   

z%TestUnitaryGroup.test_reproducibilityc                 C   r  r  )r   rJ   r   r   r  r>   r>   r?   r    r  z!TestUnitaryGroup.test_invalid_dimc                 C   r  )Nrc   r  r   r   r	   r{   )r   r   r   r  r>   r>   r?   r    r  z#TestUnitaryGroup.test_frozen_matrixc                 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  r>   r>   r?   r  	  r  z3TestUnitaryGroup.test_unitarity.<locals>.<listcomp>rD   r  r   r  rf  )r   r   r7   r   conjr   rM   r9   )rR   r9  rX   r>   r>   r?   test_unitarity	  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  r>   )r   r   eigvalsr  r>   r>   r?   r  	  r  z.TestUnitaryGroup.test_haar.<locals>.<listcomp>rD   r  )r7   r   r   r   r   r3  rA  imagrealr   ravelr   rz  r   r   rB  )rR   r   r  r9  r{  rX   r:   r>   r>   r?   r  	  s    zTestUnitaryGroup.test_haarN)r   r   r   r  r  r  rN  r  r>   r>   r>   r?   rL    s    	rL  c                   @   s4  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=d> Zejd6g d?d@dA Zejd6g dBdCdD ZdqdFdGZdHdI Zejd6g dJejdKg dLejdMdEdNgdOdP Zejj dQdR Z!dSdT Z"dUdV Z#dWdX Z$ejd6dYdZd[ Z%d\d] Z&ejd8g d^d_d` Z'ejdadbe(dddcddfdeedfdgddbggdgdhfgdidj Z)ejdkg dldmdn Z*dodp Z+d%S )rTestMultivariateTr	   rD   rb   rE   r`   r   )
|	ߌ?v
2R?rT  g0$H?rU  rU  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?)rd   r	   2   )r   r1  r  )r1  r   r  )r  r  r   ro  )gJ#R<g^<gL[<gN=<g`$<zx, loc, shape, df, ansc                 C   s(   t |||dd}||}t|| d S Nr   r   )r"   r   r   )rR   rX   r8  r9   rO  ansr   valr>   r>   r?   test_pdf_correctnessh	  s   
z&TestMultivariateT.test_pdf_correctnessc           	      C   s8   t |||dd}||}||}tt|| d S rW  )r"   r   r   r   r7   r   )	rR   rX   r8  r9   rO  rX  r   rt  ru  r>   r>   r?   test_logpdf_correctn	  s   

z%TestMultivariateT.test_logpdf_correctc                 C   s.   g d}t j|dd}t|}t|| d S )N)
r  rc   rb   r	   r  r   r\  rd   rE   r	   rO  )r"   r   r#   r   )rR   rX   rY  rX  r>   r>   r?   test_mvt_with_df_one_is_cauchyv	  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	   順 rO  r   r{   rg   rE   r  rd   *   )r"   r   r$   r  )rR   	P_VAL_MINr   r  r   ry  r>   r>   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	   rc   r`  r]  rD   )r"   r7   r   r   r   r   
call_countr   )rR   mockr   r>   r>   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 )Nrb         @rc   r	   r{   )
r7   r   rM   r"   r   isscalarr   r   r9   r   )	rR   r   r8  r9   rO  rX   r:   	n_samplesr|   r>   r>   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   r8  r9   rO  )rR   r   r>   r>   r?   test_default_arguments	  s   z(TestMultivariateT.test_default_arguments)NNNr   r	   r	   )NNrc   r   r	   rc   Nrc   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r8  r9   rO  rk  rR   r8  r9   rO  loc_ans	shape_ansdf_ansr   r>   r>   r?   test_default_args	  s   z#TestMultivariateT.test_default_argsrd   c                 C   s4   t |||}t|j| t|j| t|j| d S r6   rk  rn  r>   r>   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	   )r8  rD   rE   )r8  r9   rd   rm  r   )r   rJ   r"   dictr7   r   rM   )rR   r8  r9   rO  r>   r>   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 )Nrb   rE   r{   rD   r`  r  )r7   r   r}  r   rM   r"   r   r   )rR   r   r8  r9   r   dist2samples1samples2r>   r>   r?   r  	  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)r8  r9   rO  rV   )rt  r   r7   r   r_  r"   )rR   r<   r>   r>   r?   test_allow_singular
  s    z%TestMultivariateT.test_allow_singularr|   )r  rE   )r`   ra   rb   rE   r   )rD   rE   rb   r`   rO  r   r  c                 C   s<   t t|t||}|j|d}|j||f ksJ d S r  )r"   r7   r   rM   r   r9   )rR   r|   r   rO  r   r   r>   r>   r?   test_rvs
  s   zTestMultivariateT.test_rvsc                 C   s   t d}t d}d}g dg dg dg dg}g dg dg dg dg}t g d}tj|||||d}t||d	 |  d S )
NrE   r  r~  re   r  r  r  r  r   r  )rR   r   r   rO  r  r  r  r   r>   r>   r?   r  
  s   

z TestMultivariateT.test_cdf_signs)r	   rD   r`   c                 C   s   |  | d S r6   cdf_against_mvn_testrR   r   r>   r>   r?   $test_cdf_against_multivariate_normal
  s   z6TestMultivariateT.test_cdf_against_multivariate_normal)rE   ra   r  c                 C   s   |  dd d S )NrE   Tr{  r}  r>   r>   r?   -test_cdf_against_multivariate_normal_singular 
  s   z?TestMultivariateT.test_cdf_against_multivariate_normal_singularFc              	   C   s   t jd}d}d|jdd|d }t||||}d|jdd|d t |j|d }d|jdd||fd  | }d|jdd||fd | }	tjj	|	||d	|d
|d}
tj
j	|	||d
|d}t|
|dd d S )N   ]J>?a rE   r  rg   r	   r{   rd   rD   r@  T)rO  r  rV   r   )rV   r  Mb@?rf  )r7   r   r   r   r   signr   r,   r"   r   r   r   )rR   r   r~   r   r   rF  r   r   r  r  r:   r;   r>   r>   r?   r|  %
  s   &z&TestMultivariateT.cdf_against_mvn_testc           	      C   s   t jd}d}d}|jdt |d}d}tjj||||t j |d}tj	|||t |}tj
||t |}t||dd	 t t || d
ksPJ d S )Nr  rD   r   r  )r|   r3  rE   )r  r   r  rf  r%  )r7   r   r   r   sqrtr,   r"   r   r   tr   r   r  r   )	rR   r   r   r   rX   rO  r:   r;   	incorrectr>   r>   r?   test_cdf_against_univariate_t7
  s    z/TestMultivariateT.test_cdf_against_univariate_t)rD   rE   r`   r  r   )   no l   m, l   sO l   U iOfYl   O_ r~   Tc              	   C   s\  |r|dkrt d tj|}d|jdd|d }t||||}||}|| }||}	| d }
tjj	|	|||
|dd	}tj
d
d td|
|tj| |	| |d }W d    n1 sew   Y  t||ddd tjj	|	|||
||dd}tj
d
d td|
||| |	| |d }W d    n1 sw   Y  t||ddd d S )Nr  z4Agreement with qsimvtv is not great in singular caser  rg   rD   r{   r`   T)r   rV   r0  )invalidr  r   g-C6*?r%  r;  r'  )r  r   rV   r  )rI   r   r7   r   r   r   r   r,   r"   r   r3  r4   r   r   )rR   r   r   r~   r   rF  r   r   r  r  rO  r:   r;   r>   r>   r?   test_cdf_against_qsimvtvF
  s.   


" z*TestMultivariateT.test_cdf_against_qsimvtvc           	   	      s  d}t jd}d|jdd|d }t|||dd |||| }||}| d	 tjj| ||d
} fdd}t|||tj	j
||dd}t||jdd  fdd}t||d |d |d |d |d |d }t||d dd d S )NrE   l   	!y r  rd   r	   r{   T)r~   r`   )r   r  c                    s   t j| j S r6   )r,   r"   r   r   rW   r   rO  r   r>   r?   	integrandr
  s   zITestMultivariateT.test_cdf_against_generic_integrators.<locals>.integrand)rc  r   )qrngr%  r&  c                     s   t j| d d d  S )Nrd   )r,   r"   r   )zyxr  r>   r?   r  x
  s   r   rD   )r7   r   r   r   r   r,   r"   r   r/   qmcHaltonr   integralr1   )	rR   r   r   rF  r  r  r:   r  r;   r>   r  r?   $test_cdf_against_generic_integratorsc
  s"   

,z6TestMultivariateT.test_cdf_against_generic_integratorsc                 C   sb   t jd}t g dg dg dg}d}tj||d}|jg d|d}d	}t||d
d d S )Nl   Ja )go!@S<?^Ss@)r  gRP=@jT0@)r  r  gk`-3@gWK?r9   rO  re   r   g1%?r%  r&  )r7   r   r   r   r,   r"   r   r   )rR   r   r   rO  r   r:   r;   r>   r>   r?   test_against_matlab~
  s   
z%TestMultivariateT.test_against_matlabc                 C   s   d}t j|}|jdd}|jdd| }t d}| }|||f}t j|}t j|}	tj|d|i}
t|
|tj|g|R d|	i d S )Nl   2m rE   r{   r   r   )	r7   r   r   r   rM   r,   r"   r   r   )rR   r   r   r8  rX   r9   rO  r<   
rng_frozenrng_unfrozenr   r>   r>   r?   rH  
  s   


zTestMultivariateT.test_frozenc                    s   d}d}t jdj||fd}||j  |||f } d tjj| d} fdd}t |d	|}t||d
dd d S )Nrb   r  r  r{   r`   r   c                    s"   t d tj |  |  d S )Nr@  r   )r4   r7   r   rW   r   rO  r   r   r>   r?   _cdf_1d
  s   "z2TestMultivariateT.test_vectorized.<locals>._cdf_1drd   r  r%  r  )	r7   r   r   r   r,   r"   r   apply_along_axisr   )rR   r   r   rr   rX   r:   r  r;   r>   r  r?   test_vectorized
  s   

z!TestMultivariateT.test_vectorizedr  c                 C   sd   t jd}tjjdgdg|d   d}tj|djdg| |d}d|d  }t	||dd	 d S )
Nr  r	   r1  )r`  )r9   r   r   g-C6
?r&  )
r7   r   r   r   r   toeplitzr,   r"   r   r   )rR   r   r   rr   r:   r;   r>   r>   r?   test_against_analytical
  s
   z)TestMultivariateT.test_against_analyticalc                 C   s@   t dd}t j}tjj||d}tjd |}||ksJ d S )NrE   r  )r7   rM   r   r,   r"   r   r   )rR   r   rO  mvt_entropymvn_entropyr>   r>   r?   test_entropy_inf_df
  s
   z%TestMultivariateT.test_entropy_inf_df)r	   r  r  c                 C   s0   t jjd|d}t jj|d}t||dd d S )Nr   r  r]  r  r&  )r,   r"   r   r  r   )rR   rO  r  	t_entropyr>   r>   r?   test_entropy_1d
  s   z!TestMultivariateT.test_entropy_1dzdf, cov, ref, tolr  g}M@rx  r  r1  gNfh@r  c                 C   s.   t d}t|||}t| ||d d S )NrI  r&  )r7   r   r,   r"   r   r   )rR   rO  r   r;   tolr8  mvtr>   r>   r?   %test_entropy_vs_numerical_integration
  s   
z7TestMultivariateT.test_entropy_vs_numerical_integrationzdf, dim, ref, tol))r  r	   gDdVW?r  )r  r	   gU?r  )r[  r	   g/+?rx  )@xDr	   Z_2?r  )}Ô%ITr	   r  r  )r  r  gzg#.@r  )r  r  gD:f,@r  )r  r  1>`,@r  )r  r  r  r  )r  r  g8d?b@r  )r  r  g-YҺa@rx  )r  r  ZGra@r  )r  r  r  r  c                 C   s*   t jt||d}t| ||d d S )Nr  r&  )r,   r"   r7   rM   r   r   )rR   rO  r   r;   r  r  r>   r>   r?   test_extreme_entropy
  s   #z&TestMultivariateT.test_extreme_entropyc           
      C   s   t g dg dg dg dg dg}||j }d}tjj||d}td |d }t||d	d
 d}d}tjj||d}tjj||d}	t||	dd
 d S )N)gQ?g
ףp=
?g\(\߿(\?gGz?)gGzg{GzgQ?r'  gQ)gQ?g)\(?gQѿg)\(ܿg(\?)g      gGzgq=
ףpg\(\?g)r  g{Gzg\(\?g{GzԿg(\?r  r  r^  r  r&  i  i   rN  )r7   r   r   r,   r"   r   r   r   )
rR   _Ar   rO  mul_t_entropymul_norm_entropydf1df2	_entropy1	_entropy2r>   r>   r?   test_entropy_with_covariance  s"   
z.TestMultivariateT.test_entropy_with_covarianceF),r   r   r   	PDF_TESTSrI   r   r   rZ  r[  r^  rc  r5   rf  rj  rl  DEFAULT_ARGS_TESTSrr  r7   r   ARGS_SHAPES_TESTSrs  ru  r  ry  r   rz  r  r~  r  r|  r  r  r  r  r  rH  r  r  r  r  rM   r  r  r  r>   r>   r>   r?   rS  	  s    
(B


 **""..
$









rS  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g ddfg ddfddgdfdgdff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jg ed'ejg ed'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d0d1gfdgdgggdgdggddgd0d1ggfdgdggdggggddgd0d1gggfgd2d3 Zd4d5 Zd6d7 Zd8d9 Zd:d; Zd<d= Zd>d? Zd@dA ZdBdC ZdDdE ZdFdG ZdHdI ZdJS )KTestMultivariateHypergeomzx, m, n, expectedrE   rb   r`   r  rc   gr   r\  r	   rD   ri   iir  r  rd   ic                 C       t |||}t||dd d S )Nư>r&  )r   ri  r   rR   rX   r  r   r  valsr>   r>   r?   rn  "  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 )NrE   r	   r  r`   rb   )rX   r  r      )r   r   r   rt  r  r&  rc      )r   rr  r!   r   rs  r>   r>   r?   test_reduces_hypergeom@  s   z0TestMultivariateHypergeom.test_reduces_hypergeomc                 C   s:   t ddgdd}|jddd}t|d| d	d
 d S )NrE   r`   rb   r  r   r  r|  r   r   r9  r&  r   r   r   r   rR   r   r   r>   r>   r?   rz  K  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 )NrE   r`   r  rb   r  r  r  rD   r|  r   r   r9  r&  r  r  r>   r>   r?   test_rvs_broadcastingR  s   z/TestMultivariateHypergeom.test_rvs_broadcastingzm, n)r   r   r  r   r   )r   r   r   r   r   c                 C   s8   t ||}t|}| }|||dk< t|| d S )Nr   )r   r   r7   r8   r5  r   )rR   r  r   r:   res_exr>   r>   r?   test_rvs_gh16171W  s
   
z*TestMultivariateHypergeom.test_rvs_gh16171g zl?ro  r  gQ?gE|?g
y?r  r  )r`   ra   rc   ra   gzp?c                 C   r  NrM  r&  r   rr  r   r  r>   r>   r?   r  b  s   z"TestMultivariateHypergeom.test_pmfr  g ?r   r   c                 C   r  r  r  r  r>   r>   r?   r  v  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)rE   rc   r  r  r  )gf_?&BUп疈Cx@ؿ)r  ggs>d?ZsK)r  r  g35?r  r&  r   r   r   )rR   r  rm  r>   r>   r?   r    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 )Nrc   r  r  r  ro  r  r  g?gg(\?g(\r  r&  rb   r`   r   g333333?g333333r  r  r>   r>   r?   r    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`   rb   r  r  r   r   rt  r  r&  )r   r  r!   r   )rR   var0var1r>   r>   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`   rb   ro  r  r   r  r&  r	   rp  g̍3eY?gI{?r   )r   r  r   )rR   r  r  var2var3var4var5var6r>   r>   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`   rb   r  r  r  r   r  r&  r  ro  g      @r  )r   r   r!   r   rR   mean0r  rl  mean3r>   r>   r?   r    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 )NrE   r`   r  rb   ro  r  r   r   gUUUUUU@gUUUUUU@r  r&  )r   r   r   )rR   r  r  r>   r>   r?   r    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jg tddd}t|g  t|j	dk d S )Nre   r   r  r   r   r   r	   r   r   rD   r  )r   r   r   FFg<r&  r  r   )
r   r   r   r7   rY  r   r   intr   r9   r  r>   r>   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jg tddd}t|g  t|j	dk d S )Nre   r   r  r  gؗҜ<r&  r  rD   r  r  r  r  )
r   r  r   r   r7   rY  r   r  r   r9   )rR   r  r  r  r  r>   r>   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jg td	dd}tjg tjd	dd}t||dd t	|j
d
k d S )Nr  r	   r  r  r  r&  re   r   r  r6  )r   r   r   r   r7   r   r  r  r  r   r9   )rR   cov0r  r  r  r  r  r>   r>   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d	}t||}t||t||| t||t||| t|	 t	|| t|
 t
|| d S )
Nr   r  )rc   r  r     r  r  r  r  r  r  )r7   r   r   r8   r  r   r   rr  ri  r  r   )rR   r   r  rX   
mhg_frozenr>   r>   r?   rH    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  rb   g      @rg  r  g      %@g      /@)r   rJ   r   rr  	TypeErrorr  r>   r>   r?   test_invalid_params  s   z-TestMultivariateHypergeom.test_invalid_paramsN)r   r   r   rI   r   r   r7   r   rY  rn  r  rz  r  r  r   r  r  r  r  r  r  r  r  r  r  r  r  rH  r  r>   r>   r>   r?   r  !  sx    


"&&
	
r  c                   @   s0  e Zd Zdd Zdd Zejdddd Zd	d
 Z	ejddejdddd Z
ejdddd Zdd Zejdddd Zejdddd Zejdddd Zejddejdg dgfejdg dgfd d! Zejd"d#d$d% Zd&d' Zejddd(d) Zd*S )+TestRandomTablec                 C   s   t jdS )Nl   OfE)r7   r   r   r  r>   r>   r?   get_rng  rZ   zTestRandomTable.get_rngc                 C   s6  d}t jt|d tddggddg W d    n1 sw   Y  d}t jt|d tddgddgg W d    n1 sAw   Y  d}t jt|d tddgddg W d    n1 scw   Y  d}t jt|d tddgdd	g W d    n1 sw   Y  d
}t jt|d tddgddg W d    n1 sw   Y  d}t jt|d tddgg d W d    n1 sw   Y  d}t jt|d tddgg d W d    n1 sw   Y  ddg}g d}tddgg d\}}}t|| t|| |t|ksJ d S )Nz`row` must be one-dimensionalrB   r	   rD   z`col` must be one-dimensionalz*each element of `row` must be non-negativerd   z*each element of `col` must be non-negativerg   z'sums over `row` and `col` must be equalr   z(each element of `row` must be an integer @)r	   r	   rD   z(each element of `col` must be an integer)r  r  r	   rE   rD   r	   r	   )rI   r   rJ   r%   _process_parametersr   r7   r   )rR   rS   rowcolr  r`  r   r>   r>   r?   test_process_parameters	  sD   

z'TestRandomTable.test_process_parameterszscale,method))r	   boyett)r  	patefieldc                 C   sZ   t ddg| }t g d| }t}|j|||dd}|j||d dd}t|| d S )Nr	   rE   r  methodr   )r7   r   r%   r   r   )rR   r3  r  r  r  ctr  gotr>   r>   r?   test_process_rvs_method_on_None-  s   z/TestRandomTable.test_process_rvs_method_on_Nonec                 C   sX   ddg}g d}d}t jt|d tj||dd W d    d S 1 s%w   Y  d S )Nr	   rE   r  z$'foo' not recognized, must be one ofrB   foor  )rI   r   rJ   r%   r   )rR   r  r  rS   r>   r>   r?   $test_process_rvs_method_bad_argument9  s   "z4TestRandomTable.test_process_rvs_method_bad_argumentr  )TFr   c                    sJ  |   }ddgg d tj dd|d}|rt nt}t||r%dnd|s4 fd	d
|r<fddn}tj|ddd\}}||}	t|	t| |dd |t|d }
t	|
|	d  |
d|jdd   }||}	|	jdks}J t|	jd D ]!}t|	jd D ]}|	||f }|||f }||}t	|| qqg dg dg}t	tj|dd ||}	|	dksJ g dg dg}t	tj|dd  ||}	|	dksJ d}tjt|d |dg W d    n1 sw   Y  d}tjt|d |dgg W d    n	1 sw   Y  d}tjt|d |tjgg W d    n	1 s8w   Y  d}tjt|d |dgg W d    n	1 sXw   Y  d }tjt|d |g d!g W d    n	1 syw   Y  d"}tjt|d |ddgd#d$gg W d    d S 1 sw   Y  d S )%NrD   ra   r	   rE   rb   r  r  r|   r  r   ri  rr  c                    s   |  S r6   r>   rW   )r  original_methodr  r>   r?   r  S  rZ   z/TestRandomTable.test_pmf_logpmf.<locals>.methodc                    s   t  | S r6   )r7   r  rW   r  r>   r?   rY   U  s    z1TestRandomTable.test_pmf_logpmf.<locals>.<lambda>r   Tr   return_countsr   r&  )r  r  r	   )r   r	   r	   )rD   r	   rE   rd   r   r   )r	   rD   rD   rg   z$`x` must be at least two-dimensionalrB   z%`x` must contain only integral valuesr  z)`x` must contain only non-negative valuesz"shape of `x` must agree with `row`rH   z"shape of `x` must agree with `col`rE   rb   )r  r%   r   rm   r7   uniquer   r   r   r   r  r9   r   r   rI   r   rJ   rY  )rR   r  r   r   r   objrr  
unique_rvscountsry  p2rvs_ndr+  r,  pijrvijqijrX   rS   r>   )r  r  r  r  r?   test_pmf_logpmfB  sz   
$zTestRandomTable.test_pmf_logpmfr  )r  r  c                 C   s   |   }ddg}g d}tj||d||d}t||}tt|t| t|d|dd t|jd	d
t|d t|jdd
t|d d S )NrD   ra   r  r  r  r   r  rf  rd   r   r  rg   )r  rE   )	r  r%   r   r   r   r7   r   r   broadcast_to)rR   r  r   r  r  r   r   r>   r>   r?   test_rvs_mean  s   zTestRandomTable.test_rvs_meanc                 C   sn   |   }ddg}g d}tj||dd|d}tj||dd|d}tj|dd	}tj|dd	}t||d
d d S )NrD   ra   r  r@  r  r  r  r   r   {Gz?rf  )r  r%   r   r7   r  r   )rR   r   r  r  r  r  r  rm  r>   r>   r?   test_rvs_cov  s   zTestRandomTable.test_rvs_covc           
      C   s  ddg}g d}t j||||  d}|jdksJ t j||d||  d}|jdks-J t||d	  t j||d	||  d}|jd
ksGJ t j||d||  d}|jdksZJ t j||d||  d}|jdksmJ t|ddd|dd d}	tjt	|	d t j||d||  d W d    n1 sw   Y  tjt	|	d t j||t
j||  d W d    d S 1 sw   Y  d S )NrD   ra   r  r  r  r	   r  rH   r   )r   rD   rE   r  )r  rD   rE   )rb   r`   )rb   r`   rD   rE   rE   r  r&  z/`size` must be a non-negative integer or `None`rB   rd   )r%   r   r  r9   r   r   r  rI   r   rJ   r7   rY  )
rR   r  r  r  r   rv2rv3rv4rv5rS   r>   r>   r?   test_rvs_size  sF   
"zTestRandomTable.test_rvs_sizec           	      C   sh   ddg}g d}t }|j||d||  d}tj|ddd\}}||||}t|t| |d	d
 d S )NrD   ra   r  r_  r  r   Tr   r  r&  )r%   r   r  r7   r  rr  r   r   )	rR   r  r  r  r  r   r  r  ry  r>   r>   r?   test_rvs_method  s   zTestRandomTable.test_rvs_methodc                 C   sn   g d}g d}t ||}|jd||  d}tdt|t|f}g dg dg dg|d< t|| d S )Nr  )r	   r   r   r   r  r  )r   r   r   r   .r%   r   r  r7   r   r   r   rR   r  r  r  rc  r   r  r>   r>   r?   test_rvs_with_zeros_in_col_row  s   
z.TestRandomTable.test_rvs_with_zeros_in_col_row)Nr  r  r  r   r  c                 C   sD   t ||}|jd||  d}tdt|t|f}t|| d S )Nr  r  r  r  r>   r>   r?   test_rvs_with_edge_cases  s   
z(TestRandomTable.test_rvs_with_edge_casesr   rB  c                 C   s   dd l m  m} tjddgtjd}tjg dtjd}t|d| }t|}||||d|  }|j	dt
|t
|fksBJ t||ksKJ d S )Nr   r	   rE   r  r  	rvs_rcont)scipy.stats._rcontr,   _rcontr7   r   int64rm   r   r  r9   r   )rR   r   r  r  r  r   ntotresultr>   r>   r?   test_rvs_rcont  s   
zTestRandomTable.test_rvs_rcontc                 C   s   ddg}g d}t |||  d}| }t ||}t||  t |||}t||| t |||}t||| d S )NrD   ra   r  r   )r%   r  r   r   r   rr  ri  )rR   r  r  rc  rr  r  r>   r>   r?   rH    s   zTestRandomTable.test_frozenc                 C   sV   ddg}g d}t |||  d}t j||d||  d}|jd|d}t|| d S )NrD   ra   r  r   r  r  )r|   r  )r%   r  r   r   )rR   r  r  r  rc  r  r  r>   r>   r?   test_rvs_frozen$  s   zTestRandomTable.test_rvs_frozenN)r   r   r   r  r  rI   r   r   r  r  r  r  r  r  r  r  r  r   rH  r!  r>   r>   r>   r?   r    s:    $

	N

)


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|   ro  )r   r   pickledumpsloadsr   )distfnr<   r  r   r0	unpickledr1r>   r>   r?   check_pickling0  s   



r)  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 )NrE   r1  r   rP  r>   r   r  r`   r&  rI  )r7   rM   r   r   r   r   r   r   r   r   r3   r)  )r3  distsr%  r<   r>   r>   r?   test_random_state_propertyD  s   


	
r+  c                   @   s  e Zd Zejdg dejdg ddd Zejdddgejd	g d
dd Zdd Zdd Z	dd Z
ejd	ddgdd Zejd	ddgdd Zejdejejgdd Zejdejejgdd Zejd eg d!eg d!d"d#feg d$eg d%d"d&feg d!eg d!d'd(feg d$eg d%d'd)feg d!eed*ed+dgd,d-feg d$eg d!d,d.feg d/eg d/d,d0feg d/eed*ed+dddgd,d1feg d/eed2ed3dddgd,d4fg	d5d6 Zejd eg d!eg d!d"d7feg d$eg d%d"d8feg d!eg d!d'd9feg d$eg d%d'd:feg d!eed*ed+dgd,d;feg d$eg d!d,d<feg d/eg d/d,d=feg d/eed*ed+dddgd,d>feg d/eed2ed3dddgd,d?fg	d@dA ZejdBg dCdDdE ZejdejejgdFdG ZdHdI Zejdg dJejdKg dLdMdN ZdOdP ZdQdR ZdSdT ZdUS )VTestVonMises_Fisherr   )rD   rE   rb   ra   r|   r:  c                 C   s   t jd}t |fdt | }t|d|d}||}t |t |}}|j	|||dj
}	|j
|	ks9J t jj|dd}
t|
d d S )Nr;  r	   r   r{   rd   r   r   )r7   r   r   rP  r  r'   r   r   rM   r   r9   r   r   r   )rR   r   r|   r   r   vmf_distr  r   r   r=  r>  r>   r>   r?   r  W  s   
z TestVonMises_Fisher.test_samplesr`   ro  kappa)g  4&kCr  gꌠ9Y>)Fc                 C   s@   t jd}t |fdt | }t|||d}|d d S )Nr;  r	   r   r  )r7   r   r   rP  r  r'   r   )rR   r   r.  r   r   r-  r>   r>   r?    test_sampling_high_concentratione  s   z4TestVonMises_Fisher.test_sampling_high_concentrationc                 C   L   t d}d}tjt|d t|d W d    d S 1 sw   Y  d S )Nrp  z%'mu' must have one-dimensional shape.rB   r	   r7   rL   rI   r   rJ   r'   rR   r   ra  r>   r>   r?   test_two_dimensional_mun  
   
"z+TestVonMises_Fisher.test_two_dimensional_muc                 C   r0  )NrI  z%'mu' must be a unit vector of norm 1.rB   r	   r1  r2  r>   r>   r?   test_wrong_norm_mut  r4  z&TestVonMises_Fisher.test_wrong_norm_muc                 C   r0  )Nr  z$'mu' must have at least two entries.rB   r	   r1  r2  r>   r>   r?   test_one_entry_muz  r4  z%TestVonMises_Fisher.test_one_entry_murd   )r`   rE   c                 C   F   d}t jt|d tddg| W d    d S 1 sw   Y  d S )Nz"'kappa' must be a positive scalar.rB   r	   r   rI   r   rJ   r'   rR   r.  ra  r>   r>   r?   test_kappa_validation  s   "z)TestVonMises_Fisher.test_kappa_validationr   r   c                 C   r7  )NzFor 'kappa=0' the von Mises-Fisher distribution becomes the uniform distribution on the sphere surface. Consider using 'scipy.stats.uniform_direction' instead.rB   r	   r   r8  r9  r>   r>   r?   test_kappa_zero  s   "z#TestVonMises_Fisher.test_kappa_zeror  c                 C   sV   t g d}d}tjt|d ||ddgd W d    d S 1 s$w   Y  d S )Nr   r   r   znThe dimensionality of the last axis of 'x' must match the dimensionality of the von Mises Fisher distribution.rB   r	   r   r7   r   rI   r   rJ   rR   r  rX   ra  r>   r>   r?   test_invalid_shapes_pdf_logpdf  s
   "z2TestVonMises_Fisher.test_invalid_shapes_pdf_logpdfc                 C   sV   t ddg}d}tjt|d ||ddgd W d    d S 1 s$w   Y  d S )Nr1  r   8'x' must be unit vectors of norm 1 along last dimension.rB   r	   r   r=  r>  r>   r>   r?   test_unnormalized_input  s
   "z+TestVonMises_Fisher.test_unnormalized_inputzx, mu, kappa, reference)r   r   r   r  g0_?r<  )r   r   r   g7m0_?r  gLI/@gFK:h*7g\(\?r  i  g3pB>g7Ks@)r   r   r   r   r   go齿@gW+x\(?r%  r   gS-c                 C   "   t |||}t||dd d S )Nr  r&  )r'   r   r   )rR   rX   r   r.  	referencer   r>   r>   r?   test_pdf_accuracy     z%TestVonMises_Fisher.test_pdf_accuracyg[d6U?g1*?gJ8j#@g=fNXg][,g1VV@g|'@gmF&!gehc                 C   rB  )Nrx  r&  )r'   r   r   )rR   rX   r   r.  rC  r   r>   r>   r?   test_logpdf_accuracy  rE  z(TestVonMises_Fisher.test_logpdf_accuracyzdim, kappa, reference))rE   r  g?@)rE   r  gp&F)r`   rw  gT;ӷ&)ro  r	   goZ@c                 C   s8   t |fdt | }t|| }t||dd d S )Nr	   g+=r&  )r7   rP  r  r'   r   r   )rR   r   r.  rC  r   r   r>   r>   r?   test_entropy_accuracy  s   z)TestVonMises_Fisher.test_entropy_accuracyc                 C   s   d}t jd}tdj||d}t ddt d }d}||||}|j|ks+J t|d D ]#}t|d D ]}	||||	d d f ||}
t	|
|||	f d	d
 q9q1d S )Nrp  r;  rE   r   r  r	   r`   r   r  r&  )
r7   r   r   r&   r   rP  r  r9   r   r   )rR   r  	testshaper   rX   r   r.  
result_allr+  r,  current_valr>   r>   r?   r.    s   z%TestVonMises_Fisher.test_broadcastingc           	      C   s   t jd}t ddg}t |d |d }d}t||}t||d}tdjd|d}t |d d df |d d df }t	|
 |
  t	|||| t	|||| d S )	Nr;  r   r	   r  )r8  r.  rD   r  r   )r7   r   r   r   rA  r'   r)   r&   r   r   r   r   r   )	rR   r   r   mu_angler.  vmfvonmises_distvectorsrH  r>   r>   r?   test_vs_vonmises_2d  s   
$z'TestVonMises_Fisher.test_vs_vonmises_2d)rD   rE   ra   zkappa, mu_tol, kappa_tol))r	   r  r  )r  r9  r9  )r  {Gzt?r  )r  r%  r  c                 C   s   t |fdt | }t||}t jd}d}|j||d}	t|	\}
}t |	|
}t
|d|dd t
|||d d S )	Nr	   r;  r@  r   r   r   r  r&  )r7   rP  r  r'   r   r   r   r  arccosr   r   )rR   r   r.  mu_tol	kappa_tolr   r-  r   ri  r  mu_fit	kappa_fitangular_errorr>   r>   r?   test_fit_accuracy,  s   
z%TestVonMises_Fisher.test_fit_accuracyc                 C   L   t d}d}tjt|d t| W d    d S 1 sw   Y  d S )Nr  z'x' must be two dimensional.rB   )r7   r   rI   r   rJ   r'   r  rR   rX   ra  r>   r>   r?   #test_fit_error_one_dimensional_data=  r4  z7TestVonMises_Fisher.test_fit_error_one_dimensional_datac                 C   rX  )Nr  r@  rB   )r7   rL   rI   r   rJ   r'   r  rY  r>   r>   r?    test_fit_error_unnormalized_dataC  r4  z4TestVonMises_Fisher.test_fit_error_unnormalized_datac                 C   sf   t g d}d}t||}t||dd}|jdd}tj||dd}| }t|| t|| d S )N)r   r   r	   r`   r  r   r   )r7   r   r'   r   r   )rR   r   r.  r  r  r  r  r  r>   r>   r?   r?  I  s   

z,TestVonMises_Fisher.test_frozen_distributionN)r   r   r   rI   r   r   r  r/  r3  r5  r6  r:  r;  r'   r   r   r?  rA  r7   r   r  rD  rF  rG  r.  rO  rW  rZ  r[  r?  r>   r>   r>   r?   r,  V  s    

	





r,  c                   @   s   e Zd Zedd Zdd Zdd Zdd Zd	d
 Ze	j
dddgdd Ze	j
dddgdd Zdd Zdd Ze	j
dddgdd Ze	j
dg ddd ZdS )TestDirichletMultinomialc                 C   sJ   t jd}|jdddd}|jdd|dfd}|jdd	}|||||fS )
N   t7)r   r  rD   r{   r	   r  rd   r   )r7   r   r   r   integersr   )rR   r  r   r  rX   r   r>   r>   r?   
get_paramsX  s
   z#TestDirichletMultinomial.get_paramsc                 C   s   t jd}|ddd}|ddd}t j|dd}t||}t||t||| t|	|t	||| t|
 t
|| t| t|| t| t|| d S )Nr]  r   r  r  rd   r   )r7   r   r   r   r^  r   r(   r   ri  rr  r   r  r   )rR   r   r  rX   r   rc  r>   r>   r?   rH  `  s   
z$TestDirichletMultinomial.test_frozenc                 C   s   t g d}t |}t g d}t|||}t|||}d}t|| t|t | |j|j  kr=dks@J  J t j	
d}|ddd}|ddd}t j|d	d
}t|||}t|||}d}t|| t|t | d S )NrH   rE   rb   r`   gp.U.?r>   r]  r   r  r  rd   r   g+T<)r7   r   r   r(   rr  ri  r   r   r9   r   r   r   r^  )rR   rX   r   r  r:   logresr;   r   r>   r>   r?   test_pmf_logpmf_against_Rn  s$   

 
z2TestDirichletMultinomial.test_pmf_logpmf_against_Rc                 C   s  |  d\}}}}}|d7 }tt|||d tt|||tj  |  d\}}}}}|jdddk}t|| d ||< tt|||| d tt|||| tj  t	t||||  dksqJ t	t||||  tj ksJ d S )Nr	   r   r  r{   r1  rD   )
r_  r   r(   rr  ri  r7   r   r   roundr  )rR   r   r  r  r   rX   r+  r>   r>   r?   test_pmf_logpmf_support  s   $,z0TestDirichletMultinomial.test_pmf_logpmf_supportc                 C   s   d}dg}t |g}t||}t||d t||d d t||d t||d t j  t| | t| d t|	 d d S )Nra   r  r	   r   )
r7   r8   r(   r   rr  ri  r   r   r  r   )rR   r   r  rX   r   r>   r>   r?   test_dimensionality_one  s   
z0TestDirichletMultinomial.test_dimensionality_onemethod_namerr  ri  c                 C   sb   |  d\}}}}}tt|||}ttj|g|jR  |}||}	||jd }
t|	|
 d S Nr  r   r_  rm   r(   r,   	betabinomr   r   rR   rf  r   r  r  r   rX   r  
ref_methodr:   r;   r>   r>   r?   test_against_betabinom_pmf  s   z3TestDirichletMultinomial.test_against_betabinom_pmfr   r  c                 C   sd   |  d\}}}}}tt|||}ttj|g|jR  |}| d d df }	| }
t|	|
 d S rg  rh  rj  r>   r>   r?   test_against_betabinom_moments  s   z7TestDirichletMultinomial.test_against_betabinom_momentsc           
      C   s   t jd}d}|dd}|j|dd }t||}d}|j||d}|j|||d}t| t j|dd	d
d t|	 t j	|dd	dd | j
|	 j
  krZ|fks]J  J | }	|	j
||fksjJ t|	t |jdd tt |	|	  t tj|	d dksJ d S )Nr]  r`   r	   r  r{   r  r_  r   r   rP  r&  r9  r  )r7   r   r   r^  r(   r   r   r   r   r  r9   r   r   r   rl   r  r   r   r   )
rR   r   r   r   r  r   r  ry  rX   r   r>   r>   r?   r    s    
*"z%TestDirichletMultinomial.test_momentsc                 C   s  t g d}t |}t g d}d}tt|d tg d|| W d    n1 s/w   Y  tt|d tg d|| W d    n1 sNw   Y  d}tt|d t|g d| W d    n1 sow   Y  tt|d t|g d	| W d    n1 sw   Y  d
}tt|d t||d W d    n1 sw   Y  tt|d t||d W d    n1 sw   Y  t g d}t g d}d}tt|d t|||  W d    d S 1 sw   Y  d S )NrH   r`  z,`x` must contain only non-negative integers.rB   )r	   rd   rE   )r	   r  rE   z*`alpha` must contain only positive values.)rE   r   rb   )rE   rd   rb   z`n` must be a positive integer.ǧH@r   )r	   rD   rE   rb   z&`x` and `alpha` must be broadcastable.)r7   r   r   r   rJ   r(   ri  )rR   x0n0alpha0textrX   r  r>   r>   r?   rT     s<   
"z.TestDirichletMultinomial.test_input_validationr  c                 C   s   t g dg dg dg dg}t dgdgdgg}t g dg d	gd
}tt|}||||}|jdks<J tt|D ]2}tt|D ])}tt|D ] }||||f }	|||  ||  ||  }
t	|	|
 qRqJqBd S )Nr`  rT  r`   r`   rc   ro  r  r  ra   rc   ro  rH   )rD   rD   rE   )rD   r	   r	   rE   )rD   rE   rb   )
r7   r   r  rm   r(   r9   r   r   r   r   )rR   r  r  r   rX   r:   r+  r,  r   res_ijkr;   r>   r>   r?   test_broadcasting_pmf  s   "
$z.TestDirichletMultinomial.test_broadcasting_pmf)r   r  r   c           
      C   s   t g dg dg dg dg}t dgdgdgg}tt|}|||}|dkr1|jd	ks2J 	 tt|D ]#}tt|D ]}|||f }|||  ||  }	t||	 q@q8d S )
Nr`  rT  rr  rs  ra   rc   ro  r   )rE   rb   rE   )	r7   r   rm   r(   r9   r   r   r   r   )
rR   rf  r  r   r  r:   r,  r   rt  r;   r>   r>   r?   test_broadcasting_moments  s   "

z2TestDirichletMultinomial.test_broadcasting_momentsN)r   r   r   classmethodr_  rH  rb  rd  re  rI   r   r   rl  rm  r  rT   ru  rv  r>   r>   r>   r?   r\  W  s"    





r\  c                   @   sv   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zej	j
ej	ddd Zej	dejejejejgdd ZdS )TestNormalInverseGammac                    s   t jd}|d\}}}}t|||| tjd| |dt || |  d}t ddd}t fd	d
dt j	|fd}|
|}	t|j|	  jd|d}t|d |j\}
}|dksbJ d S )N   &( rb   rD   r	   r7  r  r`   r  c                    s     || S r6   r   rG  rX   norm_inv_gammar>   r?   rY   !      z8TestNormalInverseGamma.test_marginal_x.<locals>.<lambda>r   r<   r@  r   r   )r7   r   r   r,   normal_inverse_gammar  r  r2  r-   r   r   r   r  r   ks_1sampr   )rR   r   r   lmbdar  r  r  rX   r:   r;   r   rB  r>   r|  r?   test_marginal_x  s   &
z&TestNormalInverseGamma.test_marginal_xc                    s   t jd}|d\}}}}t|||| tj||d}t ddd}t fddt j t j|fd}|	|}	t
|j|	  jd	|d
}t|d |j\}
}|dksXJ d S )Nry  rb   rb  r   r  c                    s     | |S r6   rz  )rX   rG  r|  r>   r?   rY   6  r~  z9TestNormalInverseGamma.test_marginal_s2.<locals>.<lambda>r  r@  r   r	   )r7   r   r   r,   r  r   r2  r-   r   r   r   r  r   r  r   )rR   r   r   r  r  r  	inv_gammarG  r:   r;   r   rB  r>   r|  r?   test_marginal_s2*  s   
z'TestNormalInverseGamma.test_marginal_s2c           
      C   sx   t jd}|dd \}}}}|jddd \}}t||||||}tj||||||}	t|t |	 d S )Nry  )rb   r  r  )rD   r  r{   )	r7   r   r   r,   r  r   r   r   r  )
rR   r   r   r  r  r  rX   rG  r:   r;   r>   r>   r?   test_pdf_logpdf@  s   z&TestNormalInverseGamma.test_pdf_logpdfc           	      C   s  t jd}|d\}}}}|d\}}tt j|||||}t|t j t|d||||}t|t j t||d|||}t|t j t|||d||}t|t j t|||||d}t|d t|dt jg|||d}t|t j t|d|| }t|t jt jf t||d|	 }t|t jt jf t
jtdd t|||d  W d    d S 1 sw   Y  d S )Nry  rb   rD   rd   r   zDomain error in arguments...rB   )r7   r   r   r,   r  rY  r   r   r   r  rI   r   rJ   r   )	rR   r   r   r  r  r  rX   rG  r:   r>   r>   r?   test_invalid_and_special_casesI  s,   
"z5TestNormalInverseGamma.test_invalid_and_special_casesc                 C   s  t jd}|d}|dd }|d}|d}|d}|d}t||||}t ||||||}	dd	 |	D }
|||}|j|	d
 jksOJ t|	 tjj|
  |
||}|j|	d
 jksjJ t|	 tjj
|
  t ||||}	dd	 |	D }
| }|d
 j|	d
 jksJ t|d
 	 |d 	 ftjj|
  | }|d
 j|	d
 jksJ t|d
 	 |d 	 ftjj|
  d}t jd}|j||d}t jd}d}tjj|
||d}t|d
 ||d |f| d S )Nry  rD   )rE   r	   )rb   r	   r	   )r`   r	   r	   r	   )ra   r	   r	   r	   r	   )rc   r	   r	   r	   r	   r	   c                 S   r)  r>   r7   rR  r  arrr>   r>   r?   r  z  r*  z<TestNormalInverseGamma.test_broadcasting.<locals>.<listcomp>r   c                 S   r)  r>   r  r  r>   r>   r?   r    r*  r	   )ra   r`   rb   rE   rD   l   Mr   )ra   x   )r7   r   r   r,   r  broadcast_arraysr   r9   r   rR  r   r   r  r   r  )rR   r   r  r  r  r   rG  rX   r   broadcastedbroadcasted_raveledr:   r|   r9   r;   r>   r>   r?   r.  k  sR   









&z(TestNormalInverseGamma.test_broadcastingr  c                    s   t jd}|d\}}}}|d7 }t||||   }t fddt j t jdt j}t|d |d dd t fd	dt j t jdt j}t|d
 |d dd d S )Nry  rb   rD   c                    s     || | S r6   rz  r{  r   r>   r?   rY         z5TestNormalInverseGamma.test_moments.<locals>.<lambda>r   r  r&  c                    s     || |  S r6   rz  r{  r  r>   r?   rY     r  r	   )	r7   r   r   r,   r  r   r0   r   r   )rR   r   r   r  r  r  r:   r;   r>   r  r?   r    s     z#TestNormalInverseGamma.test_momentsr  c                 C   s  t jdk r
td t jd}|jdddd|\}}}}}}t d|}	t	
||||}
|
 d	 j|	ks:J |
 d
 j|	ksEJ |
 d	 j|	ksPJ |
 d
 j|	ks[J |
 d	 j|	ksfJ |
 d
 j|	ksqJ |
||j|	ks|J |
||j|	ksJ d S )N2z*Scalar dtypes only respected after NEP 50.ry  rE   r  ra   r{   r   r   r	   )r7   __version__rI   r   r   r   r   astyperesult_typer,   r  r   r  r   r  r   r   )rR   r  r   rX   rG  r   r  r  r  	dtype_outr   r>   r>   r?   
test_dtype  s   

"z!TestNormalInverseGamma.test_dtypeN)r   r   r   r  r  r  r  r.  rI   r   r  	fail_slowr  r   r7   int32float16float32r  r  r>   r>   r>   r?   rx    s    	"3
rx  r  )\__doc__r"  numpy.testingr   r   r   r   r   r   rI   r   r   test_continuous_basicr
   numpyr7   scipy.linalgr   scipy.stats._multivariater   r   r   scipy.statsr   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)   r*   r+   r,   scipy.integrater-   r.   r/   r0   r1   scipy.specialr2   common_testsr3   	data._mvtr4   unittest.mockr5   r@   rA   r   r   r   r  r  r>  r?  rf  r  r  r  r  r9  rL  rS  r  r  r)  r+  r,  r\  rx  r>   r>   r>   r?   <module>   s~     x	 
a      L J l - %Frg9<     e  -   =