U
    9%e[>                     @   s  d Z ddlmZ ddlmZ ddlmZ ddlmZ ddl	m
Z
mZ ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZmZ ddlmZ ddlmZmZ ddlmZmZmZ ddl m!Z!m"Z" ddl#m$Z$m%Z% ddl&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z. ddl/m0Z0 ddl1m2Z2 ddl3m4Z4 G dd de'Z5G dd dZ6G dd dZ7G dd dZ8e6e8e8e7dZ9G dd de)e(Z:G d d! d!e+Z;G d"d# d#e)Z<d$S )%zq
Joint Random Variables Module

See Also
========
sympy.stats.rv
sympy.stats.frv
sympy.stats.crv
sympy.stats.drv
    )prod)Basic)Lambda)S)DummySymbol)sympify)
ProductSetIndexed)Product)Sum	summation)Tuple)Integral	integrate)ImmutableMatrixmatrix2numpy
list2numpy)SingleContinuousDistributionSingleContinuousPSpace)SingleDiscreteDistributionSingleDiscretePSpace)ProductPSpaceNamedArgsMixinDistributionProductDomainRandomSymbolrandom_symbolsSingleDomain_symbol_converter)iterable)
filldedent)import_modulec                   @   s   e Zd ZdZdd Zedd Zedd Zedd	 Zed
d Z	edd Z
edd Zedd Zdd Zdd Zd$ddZdd Zdd Zd%d d!Zd"d# ZdS )&JointPSpacezt
    Represents a joint probability space. Represented using symbols for
    each component and a distribution.
    c                 C   s>   t |trt||S t |tr(t||S t|}t| ||S N)
isinstancer   r   r   r   r    r   __new__)clssymdist r+   S/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/sympy/stats/joint_rv.pyr'   )   s    



zJointPSpace.__new__c                 C   s   | j jS r%   )domainsetselfr+   r+   r,   r.   1   s    zJointPSpace.setc                 C   s
   | j d S )Nr   argsr/   r+   r+   r,   symbol5   s    zJointPSpace.symbolc                 C   s
   | j d S N   r1   r/   r+   r+   r,   distribution9   s    zJointPSpace.distributionc                 C   s   t | j| S r%   )JointRandomSymbolr3   r/   r+   r+   r,   value=   s    zJointPSpace.valuec                 C   s>   | j j}t|tr tt|jS t|tr8|jd d S tj	S )Nr   )
r6   r.   r&   r	   r   lenr2   r   limitsZOne)r0   _setr+   r+   r,   component_countA   s    

zJointPSpace.component_countc                    s"    fddt  jD } j| S )Nc                    s   g | ]}t  j|qS r+   r   r3   .0ir/   r+   r,   
<listcomp>L   s     z#JointPSpace.pdf.<locals>.<listcomp>)ranger=   r6   r0   r)   r+   r/   r,   pdfJ   s    zJointPSpace.pdfc                 C   s0   t | j}|st| j| jjS tdd |D  S )Nc                 S   s   g | ]}|j jqS r+   )pspacer-   r@   rvr+   r+   r,   rB   T   s     z&JointPSpace.domain.<locals>.<listcomp>)r   r6   r   r3   r.   r   r0   rvsr+   r+   r,   r-   O   s    
zJointPSpace.domainc                 C   s   | j j| S r%   )r.   r2   )r0   indexr+   r+   r,   component_domainV   s    zJointPSpace.component_domainc           
         s   j }|trtd fddt|D }dd |D }tt||}t fdd|D fdd|D }d}t|D ]<}||kr~||  j	j
j|  t|| ||< |d	7 }q~ j	jrtt j	| f| }	n$ j	jrtt j	| f| }	|	|S )
Nz_Marginal distributions cannot be computed for symbolic dimensions. It is a work under progress.c                    s   g | ]}t  j|qS r+   r>   r?   r/   r+   r,   rB   ^   s     z5JointPSpace.marginal_distribution.<locals>.<listcomp>c                 S   s   g | ]}t t|qS r+   )r   strr?   r+   r+   r,   rB   _   s     c                 3   s"   | ]}t tt j|V  qd S r%   )r   rM   r   r3   r?   r/   r+   r,   	<genexpr>a   s     z4JointPSpace.marginal_distribution.<locals>.<genexpr>c                    s   g | ]}| kr|gqS r+   r+   r?   )r)   r+   r,   rB   b   s      r   r5   )r=   Zatomsr   
ValueErrorrC   dictziptupleappendr6   r.   r2   is_Continuousr   r   is_Discreter   xreplace)
r0   indicescountorigZall_symsZreplace_dictr;   rK   rA   fr+   rD   r,   marginal_distributionY   s&    


z!JointPSpace.marginal_distributionNFc                    s   t fddtjD } p"| t fdd|D s>|S |j } D ]J}t|tr|||tt|j	|j
d i}qLt|trL|||ji}qLjt|krttdt fdd|D }t|f| S )Nc                 3   s   | ]} j | V  qd S r%   )r8   r?   r/   r+   r,   rN   p   s     z2JointPSpace.compute_expectation.<locals>.<genexpr>c                 3   s   | ]}| kV  qd S r%   r+   r?   rJ   r+   r,   rN   r   s     r5   zq
            Expectations of expression with unindexed joint random symbols
            cannot be calculated yet.c                 3   s8   | ]0}t t|j|jd   jjj|jd   fV  qdS )r5   N)r   rM   baser2   r6   r.   rG   r/   r+   r,   rN   ~   s   )rR   rC   r=   anyrE   r&   r   rV   rM   r]   r2   r   r3   r8   r   NotImplementedErrorr"   r   )r0   exprrJ   evaluatekwargssymsrH   r;   r+   )rJ   r0   r,   compute_expectationo   s     

"
zJointPSpace.compute_expectationc                 C   s
   t  d S r%   r_   r0   	conditionr+   r+   r,   where   s    zJointPSpace.wherec                 C   s
   t  d S r%   re   )r0   r`   r+   r+   r,   compute_density   s    zJointPSpace.compute_densityr+   scipyc                 C   s   t | j| | jj|||diS )zo
        Internal sample method

        Returns dictionary mapping RandomSymbol to realization value.
        )libraryseed)r   r3   r6   sample)r0   sizerk   rl   r+   r+   r,   rm      s     zJointPSpace.samplec                 C   s
   t  d S r%   re   rf   r+   r+   r,   probability   s    zJointPSpace.probability)NF)r+   rj   N)__name__
__module____qualname____doc__r'   propertyr.   r3   r6   r8   r=   rE   r-   rL   r[   rd   rh   ri   rm   ro   r+   r+   r+   r,   r$   $   s.   








	r$   c                   @   s&   e Zd ZdZdddZedd ZdS )SampleJointScipyz7Returns the sample from scipy of the given distributionNc                 C   s   |  |||S r%   )_sample_scipyr(   r*   rn   rl   r+   r+   r,   r'      s    zSampleJointScipy.__new__c           	         s   ddl }|dkst|tr*|jj|d n| ddlm  fdd fdd fddd	}d
d dd dd d	}| }|jj	|krdS ||jj	 ||}|
|||jj	 | S )zSample from SciPy.r   Nrl   )statsc                    s$   j jt| j t| j| dS )N)meancovrn   random_state)multivariate_normalrJ   r   muflattensigmar*   rn   
rand_stateZscipy_statsr+   r,   <lambda>   s
     z0SampleJointScipy._sample_scipy.<locals>.<lambda>c                    s   j jt| jt | dS )N)alpharn   r|   )	dirichletrJ   r   r   floatr   r   r   r+   r,   r      s     c                    s&   j jt| jt| jt | dS )N)nprn   r|   )multinomialrJ   intr   r   r   r   r   r   r   r+   r,   r      s
      ZMultivariateNormalDistributionZMultivariateBetaDistributionZMultinomialDistributionc                 S   s   t | j jS r%   r   r~   r   shaper*   r+   r+   r,   r          c                 S   s   t | j jS r%   r   r   r   r   r   r+   r+   r,   r      r   c                 S   s   t | j jS r%   r   r   r   r   r   r+   r+   r,   r      r   )numpyr&   r   randomdefault_rngrj   ry   keys	__class__rp   reshape)	r(   r*   rn   rl   r   Zscipy_rv_mapsample_shape	dist_listsamplesr+   r   r,   rv      s$    zSampleJointScipy._sample_scipy)N)rp   rq   rr   rs   r'   classmethodrv   r+   r+   r+   r,   ru      s   
ru   c                   @   s&   e Zd ZdZdddZedd ZdS )SampleJointNumpyz7Returns the sample from numpy of the given distributionNc                 C   s   |  |||S r%   )_sample_numpyrw   r+   r+   r,   r'      s    zSampleJointNumpy.__new__c           	         s   ddl }|dkst|tr*|jj|d n|  fdd fdd fddd}d	d d
d dd d}| }|jj|krdS ||jj |t|}|	|||jj | S )zSample from NumPy.r   Nrx   c                    s$    j t| jt t| jt|dS )N)rz   r{   rn   )r}   r   r~   r   r   r   r   r   r+   r,   r      s   
 z0SampleJointNumpy._sample_numpy.<locals>.<lambda>c                    s    j t| jt |dS )N)r   rn   )r   r   r   r   r   r   r   r+   r,   r      s    c                    s"    j t| jt| jt |dS )N)r   Zpvalsrn   )r   r   r   r   r   r   r   r   r   r+   r,   r      s     r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   )
r   r&   r   r   r   r   r   rp   r   r   )	r(   r*   rn   rl   r   Znumpy_rv_mapr   r   r   r+   r   r,   r      s"    


zSampleJointNumpy._sample_numpy)N)rp   rq   rr   rs   r'   r   r   r+   r+   r+   r,   r      s   
r   c                   @   s&   e Zd ZdZdddZedd ZdS )SampleJointPymcz6Returns the sample from pymc of the given distributionNc                 C   s   |  |||S r%   )_sample_pymcrw   r+   r+   r,   r'      s    zSampleJointPymc.__new__c           	   	      s   zddl  W n tk
r(   ddl Y nX  fdd fdd fddd}dd d	d d
d d}| }|jj|kr~dS ddl}|d|j	  
 < ||jj |  jt|dd|ddddd d }W 5 Q R X ||||jj | S )zSample from PyMC.r   Nc                    s2    j dt| jt t| jtd| jjd fdS )NXr5   r   )r~   r{   r   )ZMvNormalr   r~   r   r   r   r   r   pymcr+   r,   r      s    
 z.SampleJointPymc._sample_pymc.<locals>.<lambda>c                    s    j dt| jt dS )Nr   )a)Z	Dirichletr   r   r   r   r   r   r+   r,   r      s    c                    s.    j dt| jt| jt dt| jfdS )Nr   r5   )r   r   r   )ZMultinomialr   r   r   r   r   r   r:   r   r   r+   r,   r      s     r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   c                 S   s   t | j jS r%   r   r   r+   r+   r,   r      r   pymc3r5   F)ZdrawschainsZprogressbarZrandom_seedZreturn_inferencedataZcompute_convergence_checksr   )r   ImportErrorr   r   r   rp   logging	getLoggersetLevelERRORZModelrm   r   r   )	r(   r*   rn   rl   Zpymc_rv_mapr   r   r   r   r+   r   r,   r      s*    



0zSampleJointPymc._sample_pymc)N)rp   rq   rr   rs   r'   r   r   r+   r+   r+   r,   r      s   
r   )rj   r   r   r   c                   @   sN   e Zd ZdZdZdd Zedd Zedd Zd	d
 Z	dddZ
dd ZdS )JointDistributionz
    Represented by the random variables part of the joint distribution.
    Contains methods for PDF, CDF, sampling, marginal densities, etc.
    rE   c                 G   sN   t tt|}tt|D ]"}t|| t rt|| ||< qtj| f| S r%   )	listmapr   rC   r:   r&   r   r   r'   )r(   r2   rA   r+   r+   r,   r'     s
    zJointDistribution.__new__c                 C   s
   t | jS r%   )r   symbolsr/   r+   r+   r,   r-   %  s    zJointDistribution.domainc                 C   s   | j jd S r4   )densityr2   r/   r+   r+   r,   rE   )  s    zJointDistribution.pdfc                 C   s   t |tstd|t|f | }| jjj}| t	dd | j
D }tt|D ]^}|| jrt||| || j|||  f}qV|| jrVt||| || j|||  f}qV|S )Nz!%s should be of type dict, got %sc                 s   s   | ]}|j d  V  qdS r   Nr1   r?   r+   r+   r,   rN   2  s     z(JointDistribution.cdf.<locals>.<genexpr>)r&   rP   rO   typer   r-   r.   ZsetsrE   rR   r   rC   r:   rT   r   infrU   r   )r0   otherrJ   r<   r`   rA   r   r+   r+   r,   cdf-  s    





zJointDistribution.cdfr+   rj   Nc                 C   sf   d}||krt dt| t|s0td| t| | ||d}|dk	rN|S t d| jj|f dS )z, A random realization from the distribution )rj   r   r   r   z&Sampling from %s is not supported yet.zFailed to import %srx   Nz4Sampling for %s is not currently implemented from %s)r_   rM   r#   rO   _get_sample_class_jrvr   rp   )r0   rn   rk   rl   	librariesZsampsr+   r+   r,   rm   <  s    
zJointDistribution.samplec                 G   s
   | j | S r%   r   r0   r2   r+   r+   r,   __call__O  s    zJointDistribution.__call__)r+   rj   N)rp   rq   rr   rs   Z	_argnamesr'   rt   r-   rE   r   rm   r   r+   r+   r+   r,   r     s   


r   c                   @   s   e Zd ZdZdd ZdS )r7   zg
    Representation of random symbols with joint probability distributions
    to allow indexing."
    c                 C   sD   t | jtr@| jj|kdkr6td| j| jjd f t| |S d S )NTz$Index keys for %s can only up to %s.r5   )r&   rF   r$   r=   rO   namer   )r0   keyr+   r+   r,   __getitem__W  s    zJointRandomSymbol.__getitem__N)rp   rq   rr   rs   r   r+   r+   r+   r,   r7   R  s   r7   c                   @   sX   e Zd ZdZdd Zdd Zedd Zedd	 Zd
d Z	dd Z
dd Zdd ZdS )MarginalDistributionz
    Represents the marginal distribution of a joint probability space.

    Initialised using a probability distribution and random variables(or
    their indexed components) which should be a part of the resultant
    distribution.
    c                 G   s   t |dkr$t|d r$t|d }tdd |D sBttdtdd |D }t|t	stt t
|dkrt|S t| ||S )Nr5   r   c                 s   s   | ]}t |ttfV  qd S r%   )r&   r   r   rG   r+   r+   r,   rN   l  s     z/MarginalDistribution.__new__.<locals>.<genexpr>zMarginal distribution can be
             intitialised only in terms of random variables or indexed random
             variablesc                 s   s   | ]
}|V  qd S r%   r+   rG   r+   r+   r,   rN   p  s     )r:   r!   rR   allrO   r"   r   Zfromiterr&   r   r   r   r'   )r(   r*   rJ   r+   r+   r,   r'   i  s    zMarginalDistribution.__new__c                 C   s   d S r%   r+   r/   r+   r+   r,   checku  s    zMarginalDistribution.checkc                 C   s&   dd | j d D }tdd |D  S )Nc                 S   s   g | ]}t |tr|qS r+   )r&   r   r?   r+   r+   r,   rB   z  s     
 z,MarginalDistribution.set.<locals>.<listcomp>r5   c                 S   s   g | ]}|j jqS r+   )rF   r.   rG   r+   r+   r,   rB   {  s     )r2   r	   rI   r+   r+   r,   r.   x  s    zMarginalDistribution.setc                 C   s   | j d }dd |D S )Nr5   c                 S   s   h | ]}|j jqS r+   )rF   r3   rG   r+   r+   r,   	<setcomp>  s     z/MarginalDistribution.symbols.<locals>.<setcomp>r1   rI   r+   r+   r,   r   }  s    
zMarginalDistribution.symbolsc                    s   | j d | j d  }  fddt|D }t|trpt|jj }tdddtfdd	|D }||}ntd
d	  D }t	|| 
|| S )Nr   r5   c                    s   g | ]}| kr|qS r+   r+   r?   r\   r+   r,   rB     s      z,MarginalDistribution.pdf.<locals>.<listcomp>xT)realc                 3   s   | ]}t  |V  qd S r%   r
   r?   )r   r+   r,   rN     s     z+MarginalDistribution.pdf.<locals>.<genexpr>c                 s   s*   | ]"}t |tr|jjn|jd  V  qdS r   )r&   r   rF   r3   r2   rG   r+   r+   r,   rN     s     )r2   r   r&   r   r:   r-   r   rR   rE   r   compute_pdf)r0   r   r`   marginalise_outrX   rc   r+   )rJ   r   r,   rE     s    
zMarginalDistribution.pdfc                 C   s4   |D ]*}d}t |tr|jj}| || |}q|S r4   )r&   r   rF   rE   r   )r0   r`   rJ   rH   Zlpdfr+   r+   r,   r     s    
z MarginalDistribution.compute_pdfc                 C   s   ddl m} t|tr |jj}n$t|trD|j|j|j	d }|
||jji}|jjrrt||jj|f}n:|jjr|tjtjtjfkr|j|jf}|||jj|f}|S )Nr   )r   r5   )sympy.concrete.summationsr   r&   r   rF   r.   r   r]   rL   r2   rV   r3   rT   r   rU   r   ZIntegersZNaturalsZ	Naturals0r   sup)r0   r`   rH   r   domr+   r+   r,   r     s    


z$MarginalDistribution.marginalise_outc                 G   s
   | j | S r%   r   r   r+   r+   r,   r     s    zMarginalDistribution.__call__N)rp   rq   rr   rs   r'   r   rt   r.   r   rE   r   r   r   r+   r+   r+   r,   r   `  s   

r   N)=rs   mathr   Zsympy.core.basicr   Zsympy.core.functionr   Zsympy.core.singletonr   Zsympy.core.symbolr   r   Zsympy.core.sympifyr   Zsympy.sets.setsr	   Zsympy.tensor.indexedr   Zsympy.concrete.productsr   r   r   r   Zsympy.core.containersr   Zsympy.integrals.integralsr   r   Zsympy.matricesr   r   r   Zsympy.stats.crvr   r   Zsympy.stats.drvr   r   Zsympy.stats.rvr   r   r   r   r   r   r   r    Zsympy.utilities.iterablesr!   Zsympy.utilities.miscr"   Zsympy.externalr#   r$   ru   r   r   r   r   r7   r   r+   r+   r+   r,   <module>   s>   
(q''-<