o
    k)hM                     @   s  d dl Z d dlmZ d dlZd dlZd dlmZmZm	Z	m
Z
 d dlZd dlmZmZmZmZmZmZmZmZmZ d dlmZ zd dlZdZW n eyQ   dZY nw zd dlZdZW n eye   dZY nw ejjdkrndZe j e j!e"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" Z0d#d$ Z1G d%d& d&Z2G d'd( d(e2Z3G d)d* d*e2Z4G d+d, d,e2Z5G d-d. d.e2Z6G d/d0 d0e2Z7G d1d2 d2Z8dS )3    N)join)assert_equalassert_allcloseassert_array_equalassert_raises)		GeneratorMT19937PCG64	PCG64DXSMPhiloxRandomStateSeedSequenceSFC64default_rng)	interfaceFT   c                 C   sh   | D ]/}t | | trt| | ||  qt | | tjr't| | ||  q| | || ks1J qd S N)
isinstancedictassert_state_equalnpZndarrayr   )actualtargetkey r   >lib/python3.10/site-packages/numpy/random/tests/test_direct.pyr   '   s   r   c                 C   s   | t d? d t jS )N   g      p>)r   uint32Zastypefloat32)ur   r   r   uint32_to_float321   s   r    c                 C   s\   t | } t j| t d? t jd}t d}t j| |@ t jd}t ||g }t|S )N    dtype    )r   uint64arrayr   Zcolumn_stackZravelr    )xupperlowerZjoinedr   r   r   uniform32_from_uint645   s   

r*   c                 C   s0   t | t d? } t | t d@ } t| S )N   r$   )r   r%   r   r    r'   r   r   r   uniform32_from_uint53>   s   r-   c                 C   s   t | S r   )r    r,   r   r   r   uniform32_from_uint32D   s   r.   c                 C   s4   |dkrt | S |dkrt| S |dkrt| S t)N@   5   r!   )r*   r-   r.   NotImplementedErrorr'   bitsr   r   r   uniform32_from_uintH   s   r4   c                 C   s$   |dv rt | S |dkrt| S d S )N)r/   ?   r0   r!   )uniform_from_uint64uniform_from_uint32r2   r   r   r   uniform_from_uintS   s
   r8   c                 C   s   | t d? d S )N   g      <)r   r%   r,   r   r   r   r6   Z   s   r6   c                 C   s`   t t| d }tdt| dD ]}| | d? }| |d  d? }|d | d ||d < q|S )N   r      r      g      Ag      @C)r   emptylenrange)r'   outiabr   r   r   r7   ^   s   r7   c                 C   s   |  tjd S )N      ?)viewr   doubler,   r   r   r   uniform_from_dsfmtg   s   rG   c           
      C   s   |dv r	t | }n|dkrt| }nt| }g }d}d }}t||k rtd}|dks.|dkrTd||  d }d||d   d }|| ||  }|d7 }|dks.|dks.td	t| | }	||	|  ||	|  t||k s$|d | S )
N)r/   r5   r!   r           r:   rD   g       @r   g       )r6   r7   rG   r>   r   Zsqrtlogappend)
r'   nr3   ZdoublesgausslocZx1Zx2Zr2fr   r   r   gauss_from_uintk   s*   

rO   c                  C   s   ddl m} m}m} ttdddd}|d tdi |j}t|j|j t|j	|j	 t
t|  t
t| | }t
t|jd t|ddksMJ d S )Nr   )ISeedSequenceISpawnableSeedSequenceSeedlessSeedSequence
   )r   r:   r<   )	spawn_keyZ	pool_sizer   )numpy.random.bit_generatorrP   rQ   rR   r   r?   spawnstater   n_children_spawnedr   	TypeErrorr1   generate_stater>   )rP   rQ   rR   s1s2Zdummyr   r   r   test_seedsequence   s   


r]   c                     s   t j } | jj  d} fddtdD }dd |D |ks$J | jd} fddtddD }dd |D |ksAJ | d} fddtdd	D }d
d |D }||ks_J |d  |d  ksmJ dS )z? Test spawning new generators and bit_generators directly.
    r;   c                       g | ]} j |f qS r   rT   .0rA   seqr   r   
<listcomp>       z+test_generator_spawning.<locals>.<listcomp>c                 S   s   g | ]}|j qS r   r_   )ra   cr   r   r   rd      s    c                    r^   r   r_   r`   rb   r   r   rd      re   rS   c                 S   s   g | ]}|j jqS r   )seed_seqrT   )ra   bgr   r   r   rd      s    c                    r^   r   r_   r`   rb   r   r   rd      re      c                 S   s   g | ]}|j jjqS r   )bit_generatorrg   rT   )ra   rngr   r   r   rd      s    r   r   N)r   randomr   rj   rg   rV   r?   Zuniform)rk   Znew_ssZexpected_keysZnew_bgsZnew_rngsZ
found_keysr   rb   r   test_generator_spawning   s   


 rm   c                  C   s   ddl m}  G dd d}| | tj| }tjtdd |	d W d    n1 s0w   Y  tjtdd |j
	d W d    d S 1 sNw   Y  d S )Nr   )rP   c                   @   s   e Zd ZejfddZdS )z,test_non_spawnable.<locals>.FakeSeedSequencec                 S   s   t j||dS )Nr"   )r   Zzeros)selfZn_wordsr#   r   r   r   rZ      s   z;test_non_spawnable.<locals>.FakeSeedSequence.generate_stateN)__name__
__module____qualname__r   r   rZ   r   r   r   r   FakeSeedSequence   s    rr   zThe underlying SeedSequence)matchr;   )rU   rP   registerr   rl   r   pytestraisesrY   rV   rj   )rP   rr   rk   r   r   r   test_non_spawnable   s   
"rw   c                   @   s   e Zd ZejZi  ZZedd Z	e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ejjed!d"d#d$ Zejjed%d"d&d' Zd(d) Z d*S )+Basec                 C   s*   t | _d| _tj| _t| _g | _g | _	d S )Nr/   )
r	   rj   r3   r   r%   r#   rY   seed_error_typeinvalid_init_typesinvalid_init_valuesclsr   r   r   setup_class   s   
zBase.setup_classc                 C   s   t |=}| }|d}dd |dd  D }g }|D ]}|t|dd  d q|tj|| jddW  d    S 1 sDw   Y  d S )	N,c                 S   s   g | ]	}t | d qS )r   )intstrip)ra   sr   r   r   rd      s    z"Base._read_csv.<locals>.<listcomp>r   r   r"   )seeddata)	openreadlinesplitrJ   r   r   r   r&   r#   )r}   filenameZcsvr   r   liner   r   r   	_read_csv   s   

 $zBase._read_csvc                 C   s   | j | jd  }|d}t|| jd  | j | jd  }| }t|| jd d  | j | jd  }|d}t|| jd  d S )Nr     r   r   )rj   data1
random_rawr   data2rn   rj   Zuintsr   r   r   test_raw   s   

zBase.test_rawc                 C   sF   | j | jd  }|jdd}|d u sJ |jddd}|d u s!J d S )Nr   F)outputr   )rj   r   r   r   r   r   r   test_random_raw   s
   zBase.test_random_rawc                 C   sx   d}t | j| jd  }||}t|t| jd || j t | j| jd  }|d}t|t| jd || j d S )N   r   r   )r   rj   r   standard_normalr   rO   r3   r   )rn   rK   rsrL   r   r   r   test_gauss_inv   s   

zBase.test_gauss_invc                 C   s   t | j| jd  }t| jd | j}|t|}t|| t|j	t
j t | j| jd  }t| jd | j}|t|}t|| t|j	t
j d S )Nr   r   )r   rj   r   r8   r3   rl   r>   r   r   r#   r   Zfloat64r   rn   r   ZvalsZuniformsr   r   r   test_uniform_double   s   

zBase.test_uniform_doublec                 C   s   t | j| jd  }t| jd | j}|jt|tjd}t	|| t
|jtj t | j| jd  }t| jd | j}|jt|tjd}t	|| t
|jtj d S )Nr   r   r"   )r   rj   r   r4   r3   rl   r>   r   r   r   r   r#   r   r   r   r   r   test_uniform_float  s   

zBase.test_uniform_floatc                 C   sL   t | j| jd  }dt|v sJ t|d ddt|v s$J d S Nr   r   z#xXr'   )r   rj   r   repridr(   replacern   r   r   r   r   	test_repr  s   (zBase.test_reprc                 C   sd   t | j| jd  }dt|v sJ t| jjt|v sJ t|d ddt|vs0J d S r   )r   rj   r   strro   r   r(   r   r   r   r   r   test_str  s   (zBase.test_strc           	      C   s   dd l }| j| jd  }|j}||}||}|j}tt|dt|d ||us1J t	|| t
d}|||}t|j|j d S )Nr   r   r   d   )picklerj   r   rW   dumpsloadsr   r   r   r   r   r   )	rn   r   rj   rW   Z
bitgen_pklZreloadedZreloaded_statessZaar   r   r   test_pickle  s   


zBase.test_picklec                 C   s   dd l }| j| jd  }|j}|||}|j}t|j|j t|j|j |j	d |||}|j}t|j|j t|j
|j
 d S )Nr   r   rS   )r   rj   r   rg   r   r   r   rW   ZpoolrV   rX   )rn   r   rj   r   Zbg_plkZss_plkr   r   r   #test_pickle_preserves_seed_sequence-  s   z(Base.test_pickle_preserves_seed_sequencec                 C   sH   | j | jd  }tt dh|_W d    d S 1 sw   Y  d S )Nr   1)rj   r   ru   rv   rY   rW   rn   rj   r   r   r   test_invalid_state_type?  s   
"zBase.test_invalid_state_typec                 C   sT   | j | jd  }|j}d|d< tt ||_W d    d S 1 s#w   Y  d S )Nr   ZotherBitGeneratorrj   )rj   r   rW   ru   rv   
ValueError)rn   rj   rW   r   r   r   test_invalid_state_valueD  s   "zBase.test_invalid_state_valuec              	   C   sH   | j }| jD ]}tt ||  W d    n1 sw   Y  qd S r   )rj   rz   ru   rv   rY   rn   rj   str   r   r   test_invalid_init_typeK  s   

zBase.test_invalid_init_typec              	   C   sL   | j }| jD ]}tttf ||  W d    n1 sw   Y  qd S r   )rj   r{   ru   rv   r   OverflowErrorr   r   r   r   test_invalid_init_valuesQ  s   

zBase.test_invalid_init_valuesc                 C   sb   | j | jd  }|d |dd tt |dd W d    d S 1 s*w   Y  d S )Nr   r   rF   Zint32)rj   r   Z
_benchmarkru   rv   r   r   r   r   r   test_benchmarkW  s   
"zBase.test_benchmarkzcffi not available)reasonc                 C   :   | j | jd  }|j}t|tsJ |j}||u sJ d S Nr   )rj   r   cffir   r   )rn   rj   Zcffi_interfaceZother_cffi_interfacer   r   r   	test_cffi^  
   zBase.test_cffizctypes not availablec                 C   r   r   )rj   r   ctypesr   r   )rn   rj   Zctypes_interfaceZother_ctypes_interfacer   r   r   test_ctypesf  r   zBase.test_ctypesc                 C   sP   | j | jd  }|j}| }t|tsJ t||d  t|d ts&J d S )Nr   r   r   )rj   r   rW   __getstate__r   tupler   r   )rn   rj   rW   Z	alt_stater   r   r   test_getstaten  s   zBase.test_getstateN)!ro   rp   rq   r   r%   r#   r   r   classmethodr~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   ru   ZmarkZskipifMISSING_CFFIr   MISSING_CTYPESr   r   r   r   r   r   rx      s4    




rx   c                   @       e Zd Zedd Zdd ZdS )
TestPhiloxc                 C   sb   t | _d| _tj| _| ttd| _	| ttd| _
t| _g | _ddd d dd d fg| _d S )	Nr/   z./data/philox-testset-1.csvz./data/philox-testset-2.csv)r   Nr   r   r:   i  r   )r   rj   r3   r   r%   r#   r   r   pwdr   r   rY   ry   rz   r{   r|   r   r   r   r~   w  s   zTestPhilox.setup_classc                 C   sF   | j | jd  }|j}| j |d d |d d d}t|j|j d S )Nr   rW   counterr   )r   r   )rj   r   rW   r   )rn   rj   rW   Zkeyedr   r   r   test_set_key  s   
zTestPhilox.test_set_keyN)ro   rp   rq   r   r~   r   r   r   r   r   r   v      
r   c                   @   (   e Zd Zedd Zdd Zdd ZdS )	TestPCG64c                 C   ^   t | _d| _tj| _| ttd| _	| ttd| _
ttf| _dd gfdg| _dg| _d S )Nr/   z./data/pcg64-testset-1.csvz./data/pcg64-testset-2.csvg	@r   Nr   )r	   rj   r3   r   r%   r#   r   r   r   r   r   r   rY   ry   rz   r{   r|   r   r   r   r~        
zTestPCG64.setup_classc                 C      t | j| jd  }|jj}d}|j| |d}||j_|jdd |  |d}||j_|jddd  |  |d}||ksIJ ||ksOJ d S Nr   l        PAOi_n rS   r:      r   rj   r   rW   advanceintegersrn   r   rW   stepZval_negZval_posZval_bigr   r   r   test_advance_symmetry     


zTestPCG64.test_advance_symmetryc                 C   sl   t | d}|j}|jd }d}|d |ksJ |tdd dD  |jd }d}|d |ks4J d S )N   %@
 rW   	   r%4=Kvu1/ c                 s       | ]}d | V  qdS r:   Nr   r`   r   r   r   	<genexpr>      z/TestPCG64.test_advange_large.<locals>.<genexpr>`   r/   r!   r+   r      r:   r   l	   @H}.|zKLfbe r   rj   rW   r   sumrn   r   ZpcgrW   Zinitial_stateZadvanced_stater   r   r   test_advange_large  s   

zTestPCG64.test_advange_largeNro   rp   rq   r   r~   r   r   r   r   r   r   r     
    

r   c                   @   r   )TestPCG64DXSMc                 C   r   )Nr/   z./data/pcg64dxsm-testset-1.csvz./data/pcg64dxsm-testset-2.csvr   r   r   )r
   rj   r3   r   r%   r#   r   r   r   r   r   r   rY   ry   rz   r{   r|   r   r   r   r~     r   zTestPCG64DXSM.setup_classc                 C   r   r   r   r   r   r   r   r     r   z#TestPCG64DXSM.test_advance_symmetryc                 C   sl   t | d}|j}|j}d}|d d |ksJ |tdd dD  |jd }d}|d |ks4J d S )Nr   r   rW   c                 s   r   r   r   r`   r   r   r   r     r   z3TestPCG64DXSM.test_advange_large.<locals>.<genexpr>r   l	   oF	[UmO1X9} r   r   r   r   r   r     s   
z TestPCG64DXSM.test_advange_largeNr   r   r   r   r   r     r   r   c                   @   r   )TestMT19937c                 C   sP   t | _d| _tj| _| ttd| _	| ttd| _
t| _g | _dg| _d S )Nr!   z./data/mt19937-testset-1.csvz./data/mt19937-testset-2.csvr   )r   rj   r3   r   r   r#   r   r   r   r   r   r   ry   rz   r{   r|   r   r   r   r~     s   zTestMT19937.setup_classc                 C   s   t t| jttjg t t| jttj g t t| jttjtj g t t| jtdtjg t t| jtjg t t| jdtjg d S )Nr   )r   rY   rj   r   r&   Zpi)rn   r   r   r   test_seed_float_array  s   z!TestMT19937.test_seed_float_arrayc                 C   s   t | j| jd  }|j}|j}|d}|d |d d |d d f}||_|d}t|| |d }||_|d}t|| d S )Nr   i   rj   rW   r   pos)r   rH   )r   rj   r   rW   r   r   )rn   r   rj   rW   Zdesiredtupr   r   r   r   test_state_tuple  s   




zTestMT19937.test_state_tupleN)ro   rp   rq   r   r~   r   r   r   r   r   r   r     s
    

r   c                   @   r   )	TestSFC64c                 C   r   )Nr/   z./data/sfc64-testset-1.csvz./data/sfc64-testset-2.csvr   r   r   )r   rj   r3   r   r%   r#   r   r   r   r   r   r   rY   ry   rz   r{   r|   r   r   r   r~      s   
zTestSFC64.setup_classc                 C   s   dd l }dd l}tjg dtjd}tjtjt	d }tj
|dd}||}||}W d    n1 s;w   Y  t|tsGJ t|jd d | d S )Nr   )l   VoTMQ l   *,!;l   aE]yg    r"   r   zsfc64_np126.pkl.gzrW   )gzipr   r   r&   r%   ospathr   abspath__file__r   r   loadr   r   r   rW   )rn   r   r   Zexpected_state	base_pathZpkl_filegzZsfcr   r   r   test_legacy_pickle  s   
zTestSFC64.test_legacy_pickleN)ro   rp   rq   r   r~   r   r   r   r   r   r     r   r   c                   @   s$   e Zd Zdd Zdd Zdd ZdS )TestDefaultRNGc                 C   s4   dddddgffD ]}t | }t|jtsJ q	d S )Nr   r   )  r   i.  )r   r   rj   r	   )rn   argsrgr   r   r   	test_seed&  s   zTestDefaultRNG.test_seedc                 C   sB   t  }t|}|j|u sJ t|}||u sJ |j|u sJ d S r   )r   r   rj   )rn   rh   r   Zrg2r   r   r   test_passthrough+  s   zTestDefaultRNG.test_passthroughc                 C   s~   t d}t|}t|jtsJ |j|ju sJ tj }t	d}tj
| tjjj}t|}|j|u s7J tj
| d S )Nr   iT )r   r   r   rj   r   Z_bit_generatorr   rl   Zget_bit_generatorr	   Zset_bit_generatorZmtrandZ_rand)rn   r   r   Z	_originalrh   r   r   r   #test_coercion_RandomState_Generator3  s   

z2TestDefaultRNG.test_coercion_RandomState_GeneratorN)ro   rp   rq   r   r   r  r   r   r   r   r   %  s    r   )9r   os.pathr   sysZnumpyr   Znumpy.testingr   r   r   r   ru   Znumpy.randomr   r   r	   r
   r   r   r   r   r   Znumpy.random._commonr   r   r   ImportErrorr   r   flagsoptimizer   dirnamer   r   r   r   r    r*   r-   r.   r4   r8   r6   r7   rG   rO   r]   rm   rw   rx   r   r   r   r   r   r   r   r   r   r   <module>   sZ    ,
		 8('$&