o
    8Va3                     @   s  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
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mZmZmZ d dlmZm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/m0Z0 d dl1m2Z2m3Z3 d dl4m5Z5 edZ6dd Z7dd Z8dd Z9dd Z:dd Z;dd Z<dd Z=dd Z>dd  Z?d!d" Z@d#d$ ZAd%d& ZBd'd( ZCd)d* ZDd+d, ZEe/d-d. ZFe5d/d0 ZGdS )1    N)
logexpcosSSymbolPowsinMatrixSymbolsincpi)assumingQ)import_module)ccode)MatrixSolve)log2exp2expm1log1p)	logaddexp
logaddexp2)cosm1)optimize	cosm1_optlog2_optexp2_opt	expm1_opt	log1p_opt
optims_c99create_expand_pow_optimization
matinv_optlogaddexp_optlogaddexp2_optoptims_numpy	sinc_optsFuncMinusOneOptim)XFAILskip)compile_link_import_stringshas_c)	may_xfailcythonc                  C   s  t d} dtd|  d  td }t|tg}|dtd|  d  ks&J |t|ks/J dtd|  d  dtd  }t|tg}|dtd|  d  d ksUJ |t|ks^J t| td }t|tg}|t| kstJ |t|ks}J t| td t| d  }t|tg}|t| tdt| d   ksJ |t|ksJ td}	t|	tg}
|
|	ksJ t| d td }t|tg}t|d	ksJ |t|ksJ d S )
Nx                     zlog2(x + 3))r   r   r   r   r   rewritestr)r,   expr1opt1expr2opt2expr3opt3expr4opt4expr5opt5expr6opt6 rB   D/usr/lib/python3/dist-packages/sympy/codegen/tests/test_rewriting.pytest_log2_opt   s0     $rD   c                  C   sj   t d} dd|   }t|tg}|dt|  ksJ |t|ks#J dd|   }|t|tgks3J d S )Nr,   r2   r0   r.   )r   r   r   r   r4   r   )r,   r6   r7   r8   rB   rB   rC   test_exp2_opt6   s   rE   c                     s  t d t d } t| tg}t | dksJ |t| ks#J dt  d }t|tg}dt  |ks;J |t|ksDJ dt  d }t|tg}dt  d |ks^J |t|ksgJ tttdd}|t||gksxJ |t| |gksJ |t||gksJ dt  t  d }t|tg}dt  t  |ksJ |t|ksJ dtd   d }	t|	tg}
dtd   |
ksJ |
t|	ksJ dt  d t d  d }t|tg}| | ksJ  fd	d
}t	|||| dk sJ t d}dt  d dt|  ddt|   }t|tg}dt  t| |ks>J |t| 
 dksLJ dt  d d }t|tg}t d t  }dt  td   }|| t
 dksJ |||fv sJ |t| 
 dksJ t|}t|tg}t|}t|}|||fv sJ |t|t 
 jsJ d S )Nr,   r2   r   r.   r/   r0   FZopportunisticc                    s   |   d S )Nr.   )subsZevalf)er,   rB   rC   eve   s   ztest_expm1_opt.<locals>.evgV瞯<y   )r   r   r   r   r   r4   r%   r   Z	count_opsabsZfactorr   Zis_zero)r6   r7   r8   r9   r:   r;   Zexpm1_opt_non_opportunisticr<   r=   r>   r?   r@   rA   rJ   rK   expr7opt7expr8opt8Ztgt8aZtgt8bZexpr9Zopt9Ztgt9aZtgt9brB   rI   rC   test_expm1_optA   s`    ,$rS   c                  C   N   t td \} }t| t| d }t|tg}|t| t| ks%J d S Nx yr0   )mapr   splitr   r   r   r   r,   rK   r6   r7   rB   rB   rC   test_expm1_two_exp_terms      rZ   c                  C   s  t d} t| d }t|tg}t| | dksJ |t|ks#J dt|  d }t|tg}dt|  |ks;J |t|ksDJ dt|  d }t|tg}dt|  d |ks^J |t|ksgJ tttdd}|t||gksxJ |t||gksJ |t||gksJ dt|  t|  d }t|tg}	dt|  t|  |	ksJ |	t|ksJ dtd|   d }
t|
tg}dtd|   |ksJ |t|
ksJ ddt|   }t|tg}d	t|  |ksJ |t|ksJ d S )
Nr,   r2   r   r.   r/   r0   FrF   rL   )r   r   r   r   r   r4   r%   r   )r,   r6   r7   r8   r9   r:   r;   Zcosm1_opt_non_opportunisticr<   r=   r>   r?   r@   rA   rB   rB   rC   test_cosm1_opt   s:   r\   c                  C   rT   rU   )rW   r   rX   r   r   r   r   rY   rB   rB   rC   test_cosm1_two_cos_terms   r[   r]   c                  C   sF   t d} t| t|  d }t|ttg}|t| t|  ks!J d S )Nr,   r0   )r   r   r   r   r   r   r   r   )r,   r6   r7   rB   rB   rC   test_expm1_cosm1_mixed   s   r^   c            	      C   s  t d} t| d }t|tg}t| | dksJ |t|ks#J td|  d }t|tg}t| td |ks=J |t|  dksJJ td|  d }t|tg}td|  | dksdJ |t|ksmJ t| d }t|tg}t|dksJ d S )Nr,   r2   r   r.   r0   z
log(x + 3))r   r   r   r   r   r4   simplifyr5   )	r,   r6   r7   r8   r9   r:   r;   r<   r=   rB   rB   rC   test_log1p_opt   s    r`   c                  C   s"  t d} d|  t| td  t| d  t|  d }t|t }|t| t|  t|  t	|  ks7J |
t
t
t|ksFJ t| td t| d  }t|t}|t| t|  kseJ |
t|ksnJ t| td td|  d  }t|t}|t| td t|   }|dksJ |
t|  dksJ d|  dtd|  d  d	td   d
t|   d
 td|  d  }t|t }	|	t| dtd|  d  d	  d
t	|   td t|   }
|
dksJ |	
t
t
t|  dks	J dtd|   d }t|t}|dt	d|    }|dks)J |
t|ks3J td|  d }t|t}|t	d|  d |fv sOJ td|  d }t|t}|tdt|   }|dksmJ |
t|  dks{J td|  d }t|t}||ksJ d S )Nr,   r0   r2   r3   r   r.   r/   r-   r1      )r   r   r   r   r   r_   r   r   r   r   r4   r   )r,   r6   r7   r8   r9   r:   r;   Zdelta3r<   r=   Zdelta4r>   r?   Zdelta5r@   rA   rO   rP   Zdelta7rQ   rR   rB   rB   rC   test_optims_c99   sF   0(
 
H@(



rb   c                  C   s"  dd } t d}| |d dksJ | |d |d  dks J | |d |d  d	ks.J | t|d d
ks:J | |d dksDJ | |d dksNJ | |d  dksYJ | |d |d  dksgJ t ddd}| || |d  dks{J dd }||d t|d  dksJ d S )Nc                 S   s   t t| tdgS )NrM   r   r   r   rI   rB   rB   rC   <lambda>  s    z5test_create_expand_pow_optimization.<locals>.<lambda>r,   rM   zx*x*x*xr0   zx*x + x*x*x*xr/   zpow(x, 5) + x*x*x*xzpow(sin(x), 4)z1.0/(x*x*x*x)z
pow(x, -5)z-x*x*x*xz-x*x + x*x*x*xiTZintegerzpow(x, i) - x*xc                 S   s   t t| tddd dgS )NrM   c                 S   s   | j S N)Zis_Function)brB   rB   rC   rd     s    zGtest_create_expand_pow_optimization.<locals>.<lambda>.<locals>.<lambda>)Zbase_reqrc   rI   rB   rB   rC   rd     s    r.   z pow(x, 3) + sin(x)*sin(x)*sin(x))r   r   )Zccr,   rg   Zcc2rB   rB   rC   #test_create_expand_pow_optimization  s   $rk   c                  C   s   t ddd} td| | }td| d}tt|/ t|d | tgt||ks+J t|d | | tgt||| ks@J W d    d S 1 sKw   Y  d S )NnTrh   Ar,   r2   )r   r	   r   r   Zfullrankr   r    r   )rl   rm   r,   rB   rB   rC   test_matsolve  s   ","ro   c                  C   st   t td \} }tt| t| }t|tg}t| || dks$J t|| | dks/J |t|ks8J d S )NrV   r   )	rW   r   rX   r   r   r   r!   r   r4   rY   rB   rB   rC   test_logaddexp_opt   s   rp   c                  C   s|   t td \} }td|  d|  td }t|tg}t| || dks(J t|| | dks3J |t|ks<J d S )NrV   r0   r   )rW   r   rX   r   r   r"   r   r4   rY   rB   rB   rC   test_logaddexp2_opt)  s   rq   c                  C   s  dd } t d}| t|| t|td| d|  td| td| | dtd|  |t| |t| i t d}| t|| ||  t|| |t||  | t|| tt|t| tt|tdt| t| dtdt|  t|| t|| i d S )Nc                 S   (   |   D ]\}}t|t|ksJ qd S ri   )itemsr   r$   dkvrB   rB   rC   check3     ztest_sinc_opts.<locals>.checkr,   r0   r.   rK   )r   r   r
   rx   r,   rK   rB   rB   rC   test_sinc_opts2  s    (
r{   c                  C   s   dd } t d}| td| d|  td|  d td| td|  t|d td t|d d  t|d t|d  i d S )Nc                 S   rr   ri   rs   r   r#   rt   rB   rB   rC   rx   J  ry   z test_optims_numpy.<locals>.checkr,   r0   r2   r.   )r   r   r   r
   r   r   r   r   )rx   r,   rB   rB   rC   test_optims_numpyI  s   88
r}   c                  C   s   dd } t td \}}| t|| t||  t|| d  td| |  t|| t||  t||  t|t| | d t	|t| i d S )Nc                 S   rr   ri   r|   rt   rB   rB   rC   rx   W  ry   z%test_optims_numpy_TODO.<locals>.checkrV   r2   r0   )
rW   r   rX   r   r   r   r
   r   r   r   rz   rB   rB   rC   test_optims_numpy_TODOU  s   X$
r~   c                  C   s  t std t std td} ddtd  ddtd   d	d
td   ddtd   td t }dt|  | }tdd }|| |d}t	dt|  | t
g}d}dtt||t||d }d}t I}	td|fd|fg|	tddd\}
}t|
d| }t|
d| }d|  k rdk sJ  J d|  k rdk sJ  J W d    d S 1 sw   Y  d S ) Nzcython not installed.zNo C compiler found.r,   r0   :   u   r.   a   r/   rM   '   r-   \      
   i      z
#include <math.h>

double func_unchanged(double x) {
    return %(unchanged)s;
}
double func_rewritten(double x) {
    return %(rewritten)s;
}
)	unchanged	rewrittenz
#cython: language_level=3
cdef extern double func_unchanged(double)
cdef extern double func_rewritten(double)
def py_unchanged(x):
    return func_unchanged(x)
def py_rewritten(x):
    return func_rewritten(x)
zfunc.cz	_func.pyxZc99)Zstd)Z	build_dirZcompile_kwargsgdy=g<6S:g}:gҶOɃ;gؗҜ<)r+   r'   r)   r   r   r   r   rG   rl   r   r   dictr   tempfileZTemporaryDirectoryr(   rN   Zpy_rewrittenZpy_unchanged)r,   Z	about_twor   Zxvalrefr   ZNUMBER_OF_DIGITSZfunc_cZfunc_pyxZfoldermodinfoZerr_rewrittenZerr_unchangedrB   rB   rC   "test_compiled_ccode_with_rewritingb  s6   L	
	

"r   )Hr   Zsympyr   r   r   r   r   r   r   r	   r
   r   Zsympy.assumptionsr   r   Zsympy.externalr   Zsympy.printing.codeprinterr   Zsympy.codegen.matrix_nodesr   Zsympy.codegen.cfunctionsr   r   r   r   Zsympy.codegen.numpy_nodesr   r   Zsympy.codegen.scipy_nodesr   Zsympy.codegen.rewritingr   r   r   r   r   r   r   r   r    r!   r"   r#   r$   r%   Zsympy.testing.pytestr&   r'   Zsympy.utilities._compilationr(   r)   Z!sympy.utilities._compilation.utilr*   r+   rD   rE   rS   rZ   r\   r]   r^   r`   rb   rk   ro   rp   rq   r{   r}   r~   r   rB   rB   rB   rC   <module>   sB    0@ B&.
		
