a
    QeJM                     @   s  d dl mZ d dl mZ d dlZd dlZd dl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 d dlmZmZ d dlmZ g dg d	fZeejed
dgg7 ZddhZdd Z ej!"de  dd Z#ej!"dedd Z$ej!"dedd Z%ej!"dedd Z&ej!"dedd Z'dd  Z(d!d" Z)d#d$ Z*d%d& Z+d'd( Z,d)d* Z-ej!"d+g d,ej!"d-ed.d/ Z.ej!"d-ed0d1 Z/d2d3 Z0ej!"d4e0 d5d6 Z1d7d8 Z2ej!"d9ed:d; Z3d<d= Z4dS )>    N)assert_allclose)stats   )check_normalizationcheck_momentcheck_mean_expectcheck_var_expectcheck_skew_expectcheck_kurt_expectcheck_entropycheck_private_entropycheck_edge_supportcheck_named_argscheck_random_state_propertycheck_picklingcheck_rvs_broadcastcheck_freezing)distdiscreteinvdistdiscrete)rv_discrete_frozen)r            )皙?皙?g333333?g?values ZzipfianZ
nhypergeomc                  c   sT   t  } tD ]D\}}|tv r4tj|||tjjdV  n|||| vfV  | | q
d S )N)Zmarks)setr   distslowpytestparammarkZslowadd)seendistnameargr   r   Y/var/www/sunrise/env/lib/python3.9/site-packages/scipy/stats/tests/test_discrete_basic.pycases_test_discrete_basic   s    r(   zdistname,arg,first_casec                 C   s\  zt t| }W n ty(   | }d} Y n0 tjd |j|ddi}t|}|j| \}}t|||| d  t	|||  t
|||| d  t|| d}t||||| d  |rXd	}	|j|j|j|j|jg}
d
dddd}|| d}t||||	|
 | dkrt| t|| t|| t|| t|||  |jjtjjkrXt||tj d S )Nsample distributioniP sizei  z cdf_ppfz oth{Gz?z
 chisquare)r      r   r      )randintZ	hypergeomZ	bernoullinchypergeom_walleniusr   )getattrr   	TypeErrornprandomseedrvsuniquecheck_cdf_ppfcheck_pmf_cdf	check_othr   check_discrete_chisquarepmflogpmfcdflogcdflogsfgetr   check_scale_docstringr   r   r   r   	__class__Z_entropyrv_discreter   )r%   r&   Z
first_casedistfnr5   suppmvalphaZlocscale_defaultsZmethsZspec_kkr   r   r'   test_discrete_basic$   sF    







rJ   zdistname,argc                 C   s   zt t| }W n ty(   | }d} Y n0 |j|ddi\}}}}t|||  t|||||  t||||  t|||||  t||||||  | dvrt||||||  t	|||d t	|||||  d d S )Nr)   momentsmvsk)Zzipf	yulesimonr   r   )
r0   r   r1   r   r   r   r   r	   r
   check_moment_frozen)r%   r&   rD   rF   rG   srI   r   r   r'   test_momentsO   s    
rP   zdist,shape_argsc                 C   s   | dv }zt t| }W n, tyB   | }d| jd| jd} Y n0 td}|j}g }g }t|D ]D}|d fd|d   }	|| }
|	t
|	|
 |d	|	d	  qd|	| |	|j t|| |||tjg d S )
N)Z	betabinomskellamrM   ZdlaplaceZnchypergeom_fisherr/   zrv_discrete(values=(z, z))r   r   )r   r   r   )r0   r   r1   xkpkr2   zerosZnumargsrangeappendfullinsertr*   r   int_)distZ
shape_argsZ
shape_onlyZdistfunclocnargsallargsZbshaperI   ZshpZ	param_valr   r   r'   test_rvs_broadcastf   s$    

r^   z	dist,argsc              	   C   s   zt t| }W n ty$   | }Y n0 tjd tjdddtjddg}|j| \}}|D ]L}t	|d | || g|j
dg|R d|i|j
d	g|R d|ig qbd S )
NM r   r   
           r[         ?)r0   r   r1   r2   r3   r4   r.   supportnptassert_array_equalppf)rZ   argsrD   re_locs_a_br[   r   r   r'   test_ppf_with_loc   s    
.rm   z
dist, argsc           	      C   sP  zt t| }W n ty$   | }Y n0 tjd tjdddtjddg}|j| \}}|D ]T}|| |d | f}|jdg|R d|i|jd	g|R d|if}t	
|| qbtjjddd
dtd
tjjddd
dg}|j| \}}|D ]T}|| |d | f}|jdg|R d|i|jd	g|R d|if}t	
|| qd S )Nr_   r`   ra   r   r   rb   rc   r[   rd   )   r   r*   )r0   r   r1   r2   r3   r4   r.   re   isfrf   rg   rT   )	rZ   ri   rD   rj   rk   rl   r[   expectedresr   r   r'   test_isf_with_loc   s(    
00rs   c           	      C   s   | j |g|R  }|dt|  }t| j|g|R  ||d  t| j| j |g|R  d g|R  ||d  t| ds| j| \}}|||k  }t| j| j |g|R  d g|R  || j |d  d S )N   z
-roundtrip:0yE>rR   z ppf-cdf-next)	r=   r2   spacingrf   rg   rh   hasattrre   inc)	rD   r&   rE   msgZcdf_suppZ	cdf_supp0rk   rl   Zsupp1r   r   r'   r7      s    &
&r7   c              	   C   s8  t | dr| j}n,t| jdg|R  d }tt||d }| j|g|R  }| j|g|R   }d\}}|dkr~d\}}t	j
||d  ||d  ||d	 t|}	|	d d
 t|	d  }
t	| j|
g|R  d d}| |d|i}t	
||	dd  | t||	|  t	||
| d d S )NrR   r+   r   rb   )绽|=rz   rQ   )h㈵>r{   r   atolrtolra   r         ?r[   )rw   rR   intrh   listrU   r=   r;   cumsumrf   r   r2   Zasarraydiffassert_equal)rD   r&   r%   indexZstartindZcdfsZpmfs_cumr}   r~   rI   Z	k_shiftedr[   rZ   r   r   r'   r8      s&    

.r8   c                 C   s    t j| | ||ddd d S )Nrz   r|   )rf   r   moment)rD   r&   rF   rI   r   r   r'   rN      s    rN   c                 C   s   t j| j|g|R  d| j|g|R   ddd tddd}t j| j|g|R  | jd| g|R  ddd | jdg|R  }t | j|d g|R  dk t | j|d g|R  dk d S )	Nrd   rz   r|   r+   Gz?   r   r   )	rf   r   sfr=   r2   linspacerp   rh   assert_)rD   r&   rE   ry   qZ	median_sfr   r   r'   r9      s    (( r9   c                 C   s>  d}| j | \}}tt|d}tt|dd }	t||	}
d}|g}g }|
D ]R}| j|g|R  }|| |d krN|| |||  |}|d| krN qqN|d |k r|| |d|  t|}t|}|d }||d< t	||\}}t
t|t|| \}}t||kd	|t|t|f  d
S )as  Perform chisquare test for random sample of a discrete distribution

    Parameters
    ----------
    distname : string
        name of distribution function
    arg : sequence
        parameters of distribution
    alpha : float
        significance level, threshold for p-value

    Returns
    -------
    result : bool
        0 if test passes, 1 if test fails

    g?ii  r   r   g+=ra   ru   z2chisquare - test for %s at arg = %s with pval = %sN)re   r   maxminrU   r=   rV   r2   arrayZ	histogramr   Z	chisquarelenrf   r   str)rD   r&   r5   rH   ry   Zwsupprk   rl   lohighZdistsupportlastZdistsuppZdistmassiicurrentZhistsuppfreqZhsuppZchisZpvalr   r   r'   r:      s<    





r:   c                 C   s   | j d urtd| j v d S )NZscale)__doc__rf   r   )rD   r   r   r'   rA   4  s    
rA   method)r;   r<   r=   r>   r   r?   rh   rp   zdistname, argsc                    s   zt t| W n ty"   Y d S 0 | dv r6ddg}nddg}dd |D }ddg} j|g|R d|i}tj| fd	dt|g||R  D d
d
d d S )N)rh   rp   r   r   r   r   c                 S   s   g | ]}|gd  qS )r   r   ).0pr   r   r'   
<listcomp>H      z+test_methods_with_lists.<locals>.<listcomp>r[   c                    s   g | ]} j | qS r   )r;   )r   rG   rZ   r   r'   r   L  r   V瞯<)r~   r}   )r0   r   r1   r;   rf   r   zip)r   r%   ri   zp2r[   resultr   r   r'   test_methods_with_lists:  s    
 r   c                 C   s>   t t| }tdd}|j|g|R  }tj}t|| d S )Nrt   )r0   r   r2   Zaranger=   nanrf   r   )r%   ri   rZ   xvalsrq   r   r   r'   test_cdf_gh13280_regressionP  s
    
r   c               	   c   s   ddhi} t  }tD ]\}}||v r&q|| ztt|}W n tyT   Y qY n0 | }t|D ]2\}}|j| 	|t  v sf|j
sqf||j|fV  qfqd S )NZnbinomn)r   r   r#   r0   r   r1   _shape_info	enumeratenamer@   Zintegrality)Zintegrality_exceptionsr$   r%   shapesrZ   
shape_infoishaper   r   r'   "cases_test_discrete_integer_shapesZ  s"    


r   zdistname, shapename, shapesc                 C   s   t t| }| }dd |D }||}t|}|| }|d }	|d }
|g|	g|
gg||< |j| \}}tt||d}|j	|g|R  }t
t|dd d f rJ tt|dd d f sJ t
t|dd d f rJ d S )Nc                 S   s   g | ]
}|j qS r   )r   )r   r   r   r   r'   r   y  r   z'test_integer_shapes.<locals>.<listcomp>r   r   rn   r   r   )r0   r   r   r   r   re   r2   roundr   r;   anyisnanall)r%   Z	shapenamer   rZ   r   Zshape_namesr   Zshapes_copyZvalid_shapeZinvalid_shapeZnew_valid_shapeabr   r;   r   r   r'   test_integer_shapest  s    

  r   c                  C   s   d} t jt| d tddj W d    n1 s60    Y  t jt| d tddj W d    n1 sr0    Y  dtj_tdd}t|tsJ t	tjd d S )Nz,'rv_discrete_frozen' object has no attribute)matchrb   r   Zherringpdf)
r    ZraisesAttributeErrorr   Zbinomr   Zlogpdf
isinstancer   delattr)messageZfrozen_binomr   r   r'   test_frozen_attributes  s    ,,r   zdistname, shapesc                 C   s   t | trtt| }n| }|j| \}}t|jddgg|R  |d |f t|jddgg|R  |d |f t|j	dg|R  |d |f d S )Nr   r   )
r   r   r0   r   re   rf   r   rh   rp   interval)r%   r   rZ   r   r   r   r   r'   test_interval  s    
&&r   c                     sJ  t jd} t | dd }| d}|t | }tj||fd t dg|d d   | d }t |dd   |d d g | d }t 	|}t dg|d d   | d }t |dd   dg | d } 
 \}}	t||d  t|	|d  t || t |d t |d t jdd	N t |t | t |t j  t |t j  W d    n1 s0    Y  t || t || t |dg|d d    t jdd	p t |t  |d
d t |t  |d
d t |t  |d
d W d    n1 st0    Y  t |d |  t |d |  t |d |  t jdd	p t |t  |d
d t |t  |d
d t |t  |d
d W d    n1 sP0    Y  t || t || t |d d |dd   t d|d  t d|	  |}
t |
| t d|  | t d|  | t d|	 t d|d  t ddd} |d  |d f}t d| | t d   t d|d |	f  d}t  | t || }t || d | }t || d | |d  }t || d | |d  d }t  | t  t | t  | t jdd||||f t |d | }t d| t dd d t  | t dd | t jdd	. t  fdd     W d    n1 s0    Y  t jd}  j!d| d}t jd}  | jdd}t|| d S )Nl   KU[ rb   r   r   ra   r   r   ignore)divider   )r}   r+   r   r   r   g      ?r   rL   )rK   c                 S   s   dS )Nr   r   r   r   r   r'   <lambda>  r   z test_rv_sample.<locals>.<lambda>c                 S   s   | d S )Nr   r   r   r   r   r'   r     r   c                    s
     | S )N)r<   r   r   r   r'   r     r   d   )r*   Zrandom_statero   )"r2   r3   Zdefault_rngsortsumr   rC   r   tolistr   re   r   r;   Zerrstater<   loginfr=   r>   r   r?   rh   rp   r   r   ZmedianmeanZstdsqrtvarr   expectZentropyr5   )rngrR   rS   Zxk_leftZxk_rightZcdf2Z	cdf2_leftZ
cdf2_rightr   r   Zsf2ZpsZint2Zmed2Zmean2Zvar2Zskew2Zkurt2Zmom3r5   Zrvs0r   r   r'   test_rv_sample  s    
$,
$$4"&& 

">r   )5Znumpy.testingZtestingrf   r   numpyr2   r    Zscipyr   Zcommon_testsr   r   r   r   r	   r
   r   r   r   r   r   r   r   r   Zscipy.stats._distr_paramsr   r   Z!scipy.stats._distn_infrastructurer   r   rC   r   r(   r"   ZparametrizerJ   rP   r^   rm   rs   r7   r8   rN   r9   r:   rA   r   r   r   r   r   r   r   r   r   r   r'   <module>   sR   @

*

$

8
	

