U
    e»{jÖ$  ã                   @   s²   d Z ddlZddlZddlZddlZddlZddlmZ ddl	m
Z
mZ ddlmZ ddlmZ dZddd	„Zd
d„ Zddd„Zdd„ Zdd„ Zdd„ Zdd„ ZG dd„ dƒZdS )z
Tests for RGB Images
é    N)Úassert_equal)ÚGaussian2DKernelÚconvolve)ÚHAS_MATPLOTLIB)Ú
lupton_rgbFc                 C   s6   ddl m} |j| ddd� |r*| |¡ | ¡  |S )z<Display an rgb image using matplotlib (useful for debugging)r   NZnearestÚlower)ÚinterpolationÚorigin)Zmatplotlib.pyplotZpyplotZimshowÚtitleÚshow)Úrgbr
   Zplt© r   új/var/www/html/bwave/python/venv/lib/python3.8/site-packages/astropy/visualization/tests/test_lupton_rgb.pyÚdisplay_rgb   s    
r   c                 C   s   |   ¡ }| |k}tj||< |S )z’
    Return image with all points above satValue set to NaN.

    Simulates saturation on an image, so we can test 'replace_saturated_pixels'
    )ÚcopyÚnpÚnan)ÚimageÚsatValueÚresultZ	saturatedr   r   r   Úsaturate#   s    
r   éd   c                 C   s   t jt j d¡d | d�S )Né
   r   )Údtype)r   ÚarrayÚrandom)r   ÚNr   r   r   Úrandom_array/   s    r   c                  C   s2   t tjƒ} t | ¡}| j|jks$t‚t| |ƒ d S ©N©r   r   Zfloat64r   Úcompute_intensityr   ÚAssertionErrorr   ©Úimage_rÚ	intensityr   r   r   Útest_compute_intensity_1_float3   s    

r%   c                  C   s2   t tjƒ} t | ¡}| j|jks$t‚t| |ƒ d S r   ©r   r   Zuint8r   r    r   r!   r   r"   r   r   r   Útest_compute_intensity_1_uint:   s    

r'   c                  C   sV   t tjƒ} t tjƒ}t tjƒ}t | ||¡}| j|jks<t‚t|| | | d ƒ d S )Ng      @r   ©r#   Úimage_gÚimage_br$   r   r   r   Útest_compute_intensity_3_floatA   s    


r+   c                  C   sV   t tjƒ} t tjƒ}t tjƒ}t | ||¡}| j|jks<t‚t|| | | d ƒ d S )Né   r&   r(   r   r   r   Útest_compute_intensity_3_uintJ   s    


r-   c                   @   s‚   e Zd Z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dd„ Zdd„ Zdd„ Zdd„ Zdd„ ZdS )ÚTestLuptonRgbzA test case for Rgbc                 C   sB  t j d¡ d\| _| _| _d\}}|| _|| _||f}t  |¡}t  |¡}t  |¡}ddgddgddgddgg}ddd	d
g}	ddddg}
ddddg}t	||	|
|ƒD ]d\}}}}|t
dd| ƒ ||d |d f< |t
dd| ƒ ||d |d f< |||d |d f< q¤dd„ }tdƒ}|||ƒ| _|||ƒ| _|||ƒ| _d S )Niè  )r   é   é   )éU   éK   é   é2   é-   é   i|  iX  i N  g      ð?g      ð¿g       @g      à¿g      @r   gš™™™™™Ù?r   é   c                 S   s*   t | |ddd�}tj dd| j¡}|| S )NÚextendT)ÚboundaryZnormalize_kernelr   é   )r   r   r   ÚnormalÚshape)r   ÚpsfZconvolvedImageZrandomImager   r   r   Úconvolve_with_noiseq   s       ÿz7TestLuptonRgb.setup_method.<locals>.convolve_with_noise)r   r   ÚseedÚmin_Ústretch_ÚQÚwidthÚheightZzerosÚzipÚpowr   r#   r)   r*   )ÚselfÚmethodrC   rD   r<   r#   r)   r*   ZpointsÚvaluesZg_rZr_iÚpÚvÚgrÚrir>   r=   r   r   r   Úsetup_methodV   s,    


""zTestLuptonRgb.setup_methodc                 C   sD   t  | j| j| j¡}| | j| j| j¡}t	r@t
|t ¡ jjd� dS )z1Test creating an RGB image using an asinh stretch©r
   N)r   ÚAsinhMappingr@   rA   rB   Úmake_rgb_imager#   r)   r*   Údisplayr   ÚsysÚ	_getframeÚf_codeÚco_name)rG   ÚasinhMapÚrgbImager   r   r   Ú
test_Asinh}   s    zTestLuptonRgb.test_Asinhc                 C   sD   t  | j| j| j¡}| | j| j| j¡}tr@t|t 	¡ j
jd� dS )zHTest creating an RGB image using an asinh stretch estimated using zscalerO   N©r   ÚAsinhZScaleMappingr#   r)   r*   rQ   rR   r   rS   rT   rU   rV   ©rG   ÚmaprX   r   r   r   Útest_AsinhZscale„   s    zTestLuptonRgb.test_AsinhZscalec                 C   sD   t  | j| j| j¡}| | j| j| j¡}tr@t|t 	¡ j
jd� dS )zk
        Test creating an RGB image using an asinh stretch estimated using zscale on the intensity
        rO   NrZ   r\   r   r   r   Útest_AsinhZscaleIntensityŒ   s    z'TestLuptonRgb.test_AsinhZscaleIntensityc                 C   sˆ   dddg}|  j |d 7  _ |  j|d 7  _|  j|d 7  _tj| j | j| j|d�}| | j | j| j¡}tr„t|t 	¡ j
jd� d	S )
z�Test creating an RGB image using an asinh stretch estimated using zscale on the intensity
        where the images each have a pedestal addedr   i�  ipþÿÿr   r7   r:   )ÚpedestalrO   N)r#   r)   r*   r   r[   rQ   rR   r   rS   rT   rU   rV   )rG   r`   r]   rX   r   r   r   Ú!test_AsinhZscaleIntensityPedestal–   s    
   ÿz/TestLuptonRgb.test_AsinhZscaleIntensityPedestalc                 C   s<   t  | j¡}| | j| j| j¡}tr8t|t ¡ jj	d� dS )zlTest creating a black-and-white image using an asinh stretch estimated
        using zscale on the intensityrO   N)
r   r[   r#   rQ   rR   r   rS   rT   rU   rV   r\   r   r   r   Útest_AsinhZscaleIntensityBW¦   s    z)TestLuptonRgb.test_AsinhZscaleIntensityBWzrequires matplotlib)Úreasonc              
   C   st   d}t jdd��Z}t| j|ƒ}t| j|ƒ}t| j|ƒ}tj|||| j| j	| j
|d� tj |j¡sft‚W 5 Q R X dS )z"Test the function that does it allç     @�@ú.png©Úsuffix)ÚfilenameN)ÚtempfileÚNamedTemporaryFiler   r#   r)   r*   r   Úmake_lupton_rgbr@   rA   rB   ÚosÚpathÚexistsÚnamer!   ©rG   r   ÚtempÚredÚgreenÚbluer   r   r   Útest_make_rgb¯   s          ÿzTestLuptonRgb.test_make_rgbc                 C   sp   t  d¡ d}tjdd��L}t| j|ƒ}t| j|ƒ}t| j|ƒ}tj	|||| j
| j| jdd|d�	 W 5 Q R X d S )Nz(saturation correction is not implementedrd   re   rf   r7   éÐ  )Zsaturated_border_widthZsaturated_pixel_valuerh   )ÚpytestÚskipri   rj   r   r#   r)   r*   r   rk   r@   rA   rB   rp   r   r   r   Útest_make_rgb_saturated_fix¼   s"    
÷z)TestLuptonRgb.test_make_rgb_saturated_fixc                 C   s<   t  dd¡}| | j| j| j¡}tr8t|t 	¡ j
jd� dS )z%Test using a specified linear stretchgfffffæ Àgáz®Gá*@rO   N)r   ÚLinearMappingrQ   r#   r)   r*   rR   r   rS   rT   rU   rV   r\   r   r   r   Útest_linearÐ   s    zTestLuptonRgb.test_linearc                 C   s>   t j| jd�}| | j| j| j¡}tr:t|t 	¡ j
jd� dS )z=Test using a min/max linear stretch determined from one image)r   rO   N)r   rz   r*   rQ   r#   r)   rR   r   rS   rT   rU   rV   r\   r   r   r   Útest_linear_min_maxØ   s    z!TestLuptonRgb.test_linear_min_maxc                 C   s  t  d¡ d}t| j|ƒ| _t| j|ƒ| _t| j|ƒ| _t | j| j| jdd¡ t 	| j 
¡  ¡ ¡ ¡ slt‚t 	| j 
¡  ¡ ¡ ¡ sˆt‚t 	| j 
¡  ¡ ¡ ¡ s¤t‚| j 
¡ | _| j 
¡ | _| j 
¡ | _t | j| j| j¡}| | j| j| j¡}t�r
t|t ¡ jjd� dS )z*Test interpolationolating saturated pixelsz8replaceSaturatedPixels is not implemented in astropy yetrd   r7   rv   rO   N)rw   rx   r   r#   r)   r*   r   ZreplaceSaturatedPixelsr   ÚisfiniteZgetImageZgetArrayÚallr!   ZimagesRZimagesGZimagesBrP   r@   rA   rB   rQ   rR   r   rS   rT   rU   rV   )rG   r   rW   rX   r   r   r   Útest_saturatedà   s,    
    ÿzTestLuptonRgb.test_saturatedc              	   C   sB   t jtdd��* | j | j| j¡}t || j	| j
¡ W 5 Q R X d S )Nzshapes must match)Úmatch)rw   ZraisesÚ
ValueErrorr#   ZreshaperD   rC   r   rk   r)   r*   )rG   r#   r   r   r   Útest_different_shapes_assertsû   s    z+TestLuptonRgb.test_different_shapes_assertsN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__rN   rY   r^   r_   ra   rb   rw   ÚmarkZskipifr   ru   ry   r{   r|   r   r‚   r   r   r   r   r.   S   s   '
	
r.   )N)r   )r†   rl   rS   ri   Únumpyr   rw   Znumpy.testingr   Zastropy.convolutionr   r   Z"astropy.utils.compat.optional_depsr   Zastropy.visualizationr   rR   r   r   r   r%   r'   r+   r-   r.   r   r   r   r   Ú<module>   s$   

		