a
    
d"                  	   @   s  d dl Zd dlZd dlmZ d dlZd dlmZ dZdZdZ	ej
dg dej
d	d
dgej
dejjej
dg ddd Zej
dddgej
d	d
dgdd Zej
dg dej
dg dej
d	d
dgej
dejejgdd Zej
dg dej
d	ddgej
dejjej
dg ddd Zej
dddgej
d	ddgdd  Zej
d	ddgej
dd!d"gej
d#g d$ej
dejejgd%d& Zej
dd'gej
d	d(d)gej
dejjej
dg dej
dg dd*d+ Zej
dddgej
d	d(d)gd,d- Zej
d#g d.ej
d	d(d)gd/d0 ZdS )1    N)assert_allclose)datagư>gvIh%<=gHz>wavelet)Zdb2Zsym4Zcoif5	transformZdwtswtmodedtype)float32float64Z	complex64
complex128c                 C   s   t  d d |}|jjdkr4|d d d j|_|dkr|dkrtt	" t
j|| ||d W d    n1 sv0    Y  d S t
j|| ||d}t|tsJ t|d tjsJ t
|}|jjjdkrtnt}t||||d	 d S )
N@   cr   periodizationr   r   r   frtolatol)r   ecgastyper   kindrealimagpytestraises
ValueErrorpywtmra
isinstancelistnpndarrayimra
tol_single
tol_doubler   r   r   r   r   xcoeffsyr    r*   \/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/pywt/tests/test_mra.pytest_mra_roundtrip   s    0
r,   zrbio1.3zbior2.4c                 C   s   t j}t d d |}t| jr,J |dkrxd}tj	t
|d  tj|| |d}W d    q1 sl0    Y  ntj|| |d}t|}|jjjdkrtnt}t||||d d S )Nr   r   z,norm=True, but the wavelet is not orthogonalmatchr   r   r   )r!   r
   r   r   r   r   Wavelet
orthogonalr   warnsUserWarningr   r#   r   r   r   r$   r%   r   r   r   r   r'   msgr(   r)   r   r*   r*   r+    test_mra_warns_on_non_orthogonal1   s    0
r6   axis)r   r         ndim)r8   r9      c                 C   s  |dkrt  d d }nV|dkr<t  d dd df }n4|dkrpt  d dd df }tj|fd dd	}|j|d
d}||j k s||jkrttj	" t
j|d| |d W d    n1 s0    Y  d S t
j|d| |d}t
|}|jjjdkrtnt}t||||d d S )Nr8   r   r9       r<   0      r   r7   Fcopydb1)r   r7   r   r   )r   r   camerar!   stackr   r;   r   r   	AxisErrorr   r   r#   r   r   r   r$   r%   r   )r   r;   r7   r   r'   r(   r)   r   r*   r*   r+   test_mra_axisF   s     0
rG   Zdwt2swt2c                 C   s   t  d dd df j|dd}|jjdkrH|d d dd d f j|_|dkr|dkrtt	" t
j|| ||d	 W d    n1 s0    Y  d S t
j|| ||d	}t|tsJ t|d
 tjsJ t
|}|jjjdkrtnt}t||||d d S )Nr=      FrA   r   r   rH   r   r   r   r   r   )r   rD   r   r   r   r   r   r   r   r   r   mra2r   r    r!   r"   imra2r$   r%   r   r&   r*   r*   r+   test_mra2_roundtripe   s    "0
rL   c                 C   s   t j}t d dd df j|dd}t| jr8J |dkrd}tj	t
|d  tj|| |d}W d    q1 sx0    Y  ntj|| |d}t|}|jjjd	krtnt}t||||d
 d S )Nr=   r?   FrA   rH   3norm=True, but the wavelets used are not orthogonalr-   r/   r   r   )r!   r
   r   rD   r   r   r0   r1   r   r2   r3   rJ   rK   r   r   r   r$   r%   r   r4   r*   r*   r+   !test_mra2_warns_on_non_orthogonal   s    "0
rN   r9   r<   axes)r   r8   r   r   r9   r:   r8   r      c                    s   t  d dd df j|dd |dkr>tj fd dd t fd	d
|D rttj" t	j
 d| |d W d    n1 s0    Y  d S t	j
 d| |d}t	|} jjjdkrtnt}t |||d d S )Nr=   rI   FrA   r<   r?   r   r@   c                    s"   g | ]}| j  k p| j kqS r*   r;   .0r7   r'   r*   r+   
<listcomp>       z"test_mra2_axes.<locals>.<listcomp>rC   r   rO   r   r   )r   rD   r   r!   rE   anyr   r   rF   r   rJ   rK   r   r   r   r$   r%   r   )r   rO   r;   r   r(   r)   r   r*   rZ   r+   test_mra2_axes   s    "0
r_   Zsym2dwtnswtnc           	      C   sr  |dkr$t  d d j|dd}nj|dkrPt  d dd df j|dd}n>|dkrt  d dd df j|dd}tj|fd d	d
}|jjdkr|d d d	df j|_	|dkr|dkrt
t" tj|| ||d W d    n1 s0    Y  d S tj|| ||d}t|ts&J t|d tjs<J t|}|jjjdkrZtnt}t||||d d S )Nr8   r>   FrA   r9   rI   r?   r<   r   r@   r   .ra   r   r   r   r   r   )r   r   r   rD   r!   rE   r   r   r   r   r   r   r   r   mranr   r    r"   imranr$   r%   r   )	r   r   r   r   r;   r'   r(   r)   r   r*   r*   r+   test_mran_roundtrip   s(    $"

0
rd   c                 C   s   t j}t d dd df j|dd}t| jr8J |dkrd}tj	t
|d  tj|| |d}W d    q1 sx0    Y  ntj|| |d}t|}|jjjd	krtnt}t||||d
 d S )Nr=   r?   FrA   ra   rM   r-   r/   r   r   )r!   r
   r   rD   r   r   r0   r1   r   r2   r3   rb   rc   r   r   r   r$   r%   r   r4   r*   r*   r+   !test_mran_warns_on_non_orthogonal   s    "0
re   )rP   rQ   rS   rT   rU   )r:   rR   r   )r   r9   r8   )r      r8   )r   )r8   )r9   )rR   )r:   )c                    s   t j}t d dd df j|dd t j fd dd}t fdd	| D rtt j	" t
j d
d| d W d    n1 s0    Y  d S t
j|d
d| d}t
|}|jjjdkrtnt}t||||d d S )Nr=   rI   FrA   r?   r   r@   c                    s"   g | ]}| j  k p| j kqS r*   rW   rX   rZ   r*   r+   r[      r\   z"test_mran_axes.<locals>.<listcomp>rC   r`   r]   r   r   )r!   r
   r   rD   r   rE   r^   r   r   rF   r   rb   rc   r   r   r   r$   r%   r   )rO   r   r   Zx3dr(   r)   r   r*   rZ   r+   test_mran_axes   s    "0
rh   )numpyr!   r   Znumpy.testingr   r   r   r$   r%   r   markZparametrizeZModesmodesr,   r6   r
   r   rG   rL   rN   r_   rd   re   rh   r*   r*   r*   r+   <module>   sj   