U
    e»{j«0  ã                   @   sö   d Z ddlZddl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 ddlmZ d	Zd
d„ Zdd„ ZG dd„ de
ƒZG d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S )(zA
Tests that relate to evaluating models with quantity parameters
é    N)Úassert_allclose)Úunits)ÚModel)Ú
Gaussian1DÚPix2Sky_TANÚScaleÚShift)Úassert_quantity_allclose)Ú
UnitsErrorzT{}: Units of input 'x', {}.*, could not be converted to required input units of {}.*c               	   C   s  t dddƒ} t dtj dtj dtj ƒ}t|dtj ƒ| dƒtj ƒ tjtt 	ddd¡d�� |dƒ W 5 Q R X t|dƒ| dƒtj ƒ t
|dtj ƒj| d	ƒƒ tjtt 	dd
d¡d�� |dtj ƒ W 5 Q R X tjtdd�� | dtj ƒ W 5 Q R X dS )zn
    Test evaluation of a single model with Quantity parameters that do
    not explicitly require units.
    é   gš™™™™™¹?r   Ú zm ©Úmatchr   gü©ñÒMb@?ç      à?ÚsÚmé   zACan only apply 'subtract' function to dimensionless quantities .*N)r   ÚuÚJr   r	   ÚpytestÚraisesr
   ÚMESSAGEÚformatr   ÚkmÚvaluer   )ÚgZgq© r   úp/var/www/html/bwave/python/venv/lib/python3.8/site-packages/astropy/modeling/tests/test_quantities_evaluation.pyÚtest_evaluate_with_quantities   s    	þr   c               	   C   s~   t dtj dtj dtj ƒ} tjtt ddd¡d�� | dtj	 ƒ W 5 Q R X t
| dtj	 d	t ¡ id
�| dtj ƒƒ dS )zN
    We now make sure that equivalencies are correctly taken into account
    r   é
   é   r   ÚPHzÚnmr   é   Úx©ZequivalenciesgŠÔ3<uü#@N)r   r   ZJyr"   r   r   r
   r   r   r!   r	   Úspectral)r   r   r   r   Ú/test_evaluate_with_quantities_and_equivalenciesE   s     ÿr'   c                   @   s   e Zd ZdZdZdd„ ZdS )ÚMyTestModelr    r   c                 C   s   t d|ƒ t d|ƒ || S )NÚaÚb)Úprint)Úselfr)   r*   r   r   r   ÚevaluateZ   s    

zMyTestModel.evaluateN)Ú__name__Ú
__module__Ú__qualname__Ún_inputsÚ	n_outputsr-   r   r   r   r   r(   V   s   r(   c                   @   sL   e Z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S )ÚTestInputUnitsc                 C   s   t ƒ | _d S ©N)r(   Úmodel)r,   Úmethodr   r   r   Úsetup_methoda   s    zTestInputUnits.setup_methodc                 C   sR   t |  dd¡dƒ t |  dtj d¡dtj ƒ t |  dtj d¡dtj ƒ d S )Nr   é   é   é   é   )r	   r5   r   r   Údeg©r,   r   r   r   Útest_evaluated   s    zTestInputUnits.test_evaluatec              	   C   sÞ   dt ji| j_t|  dt j d¡dt j ƒ t|  dt j d¡dt j ƒ t|  dt j dt j ¡dt j t j ƒ tjt	t
 ddd	¡d
�� |  dt j d¡ W 5 Q R X tjt	t
 ddd	¡d
�� |  dd¡ W 5 Q R X d S )Nr$   r   r:   é   r    é   r(   r   r<   r   r   )r   r<   r5   Ú_input_unitsr	   Úradr   r   r   r
   r   r   r=   r   r   r   Útest_input_unitsj   s    *zTestInputUnits.test_input_unitsc              	   C   sž   dt ji| j_d| j_t|  dt j d¡dt j ƒ t|  dt j d¡dt j ƒ tjt	t
 dd	d
¡d�� |  dt j d¡ W 5 Q R X t|  dd¡dƒ d S )Nr$   Tr   r:   r?   r    r@   r(   r   r<   r   é	   )r   r<   r5   rA   Ú _input_units_allow_dimensionlessr	   rB   r   r   r
   r   r   r   r=   r   r   r   Ú$test_input_units_allow_dimensionlessw   s    z3TestInputUnits.test_input_units_allow_dimensionlessc                 C   sl   dt ji| j_d| j_t|  dt j d¡dt j ƒ |  tjt j d¡}t|dt j ƒ |j	t jksht
‚d S )Nr$   Tr   r:   r?   r    ih  )r   r<   r5   rA   Z_input_units_strictr	   ÚnpÚpirB   ÚunitÚAssertionError©r,   Úresultr   r   r   Útest_input_units_strictƒ   s    z&TestInputUnits.test_input_units_strictc              	   C   sŠ   dt ji| j_tjtt ddd¡d�� |  dt j	 d¡ W 5 Q R X dt  
¡ i| j_t|  dt j	 d¡ddt j	 jt jt  
¡ d� ƒ d S )Nr$   r(   r!   Úmicronr   r   r%   )r   rN   r5   rA   r   r   r
   r   r   r!   r&   Úinput_units_equivalenciesr	   Útor=   r   r   r   Útest_input_units_equivalencies�   s     ÿþz-TestInputUnits.test_input_units_equivalenciesc                 C   sR   dt ji| j_dt ji| j_|  dt j d¡}t|dt j ƒ |jt jksNt‚d S )NÚzr   r:   r?   ©	r   r<   r5   rA   rB   Z_return_unitsr	   rI   rJ   rK   r   r   r   Útest_return_unitsœ   s
    z TestInputUnits.test_return_unitsc                 C   sN   dt ji| j_t j| j_|  dt j d¡}t|dt j ƒ |jt jksJt‚d S )Nr$   r   r:   r?   rS   rK   r   r   r   Útest_return_units_scalar¥   s
    
z'TestInputUnits.test_return_units_scalarN)r.   r/   r0   r7   r>   rC   rF   rM   rQ   rT   rU   r   r   r   r   r3   `   s   
	r3   c                  C   sz   t dtj ƒ} t dtj ƒ}| |@ }|dtj dtj ƒ}t|d dtj dtj  ƒ t|d dtj dtj  ƒ dS )ú-
    Test units to first model in chain.
    r   r;   r   r   N©r   r   r<   Z	arcsecondr	   Úarcsec©Ús1Ús2ÚcsÚoutr   r   r   Útest_and_input_units²   s    r^   c                  C   sP   t dtj ƒ} t dtj ƒ}| | }|dtj ƒ}t|dtj dtj  ƒ dS ©rV   r   r;   NrW   rY   r   r   r   Útest_plus_input_unitsÁ   s
    r`   c                  C   sP   t dtj ƒ} t dtj ƒ}| |B }|dtj ƒ}t|dtj dtj  ƒ dS r_   rW   rY   r   r   r   Útest_compound_input_unitsÏ   s
    ra   c               	   C   sZ   t dtj ƒ} t dtj ƒ}| |B }tjtt ddd¡d�� |dtj ƒ W 5 Q R X dS )z:
    Test incompatible units to first model in chain.
    r   r   Úpixr<   r   N)	r   r   r<   r   r   r
   r   r   rb   ©rZ   r[   r\   r   r   r   Útest_compound_input_units_failÝ   s
    rd   c               	   C   sZ   t dtj ƒ} t dtj ƒ}| |B }tjtt ddd¡d�� |dtj ƒ W 5 Q R X dS )z1
    Test incompatible model units in chain.
    r   r   rb   r<   r   N)	r   r   rb   r<   r   r   r
   r   r   rc   r   r   r   Ú%test_compound_incompatible_units_failê   s
    re   c                  C   s`   t dtj ƒ} t dtj ƒ}| |B }|dtj dt dtj tj ¡id�}t|dtj ƒ dS )zo
    Check that equivalencies work when passed to evaluate, for a chained model
    (which has one input).
    r   r$   r   r%   é   N©r   r   r<   rb   Úpixel_scaler	   rY   r   r   r   Útest_compound_pipe_equiv_call÷   s
    (ri   c                  C   s”   t dtj ƒ} t dtj ƒ}| |@ }|dtj dtj t dtj tj ¡t dtj tj ¡dœd�}t|d dtj ƒ t|d dtj ƒ dS )	zk
    Check that equivalencies work when passed to evaluate, for a composite model
    with two inputs.
    r   r   ©Úx0Úx1r%   r   r9   r   Nrg   rY   r   r   r   Útest_compound_and_equiv_call  s    þýrm   c               	   C   s^  t dtj ƒ} dt dtj tj ¡i| _t dtj ƒ}t dtj ƒ}| |B }|jdt dtj tj ¡iksrt‚|dtj ƒ}t|dtj ƒ || B }|jdks¦t‚|dtj ƒ}t|dtj ƒ | |@ }|jdt dtj tj ¡iksðt‚| d¡}|dtj dtj ƒ}t|dtj ƒ t	j
tt d	d
d¡d�� |dtj dtj ƒ}W 5 Q R X dS )zF
    Test setting input_units_equivalencies on one of the models.
    r   r$   r   rf   Nr;   rk   Ú	TestModelr   rb   r<   r   )r   r   r<   rh   rb   rO   rJ   r	   Úrenamer   r   r
   r   r   )rZ   r[   Úspr\   r]   r   r   r   Ú'test_compound_input_units_equivalencies  s&    $$
rq   c                  C   sæ   G dd„ dt ƒ} | dƒ}t dƒ}||B }|dtj ƒ}t|dtj ƒ |jtjksVt‚||B }|dtj ƒ}t|dtj ƒ |jtjksŒt‚||@ }|dtj dtj ƒ}t|dtj ƒ |d jtjksÎt‚|d jtjksât‚d	S )
z?
    Test setting input_units_strict on one of the models.
    c                   @   s   e Zd ZdejiZdS )z6test_compound_input_units_strict.<locals>.ScaleDegreesr$   N©r.   r/   r0   r   r<   Úinput_unitsr   r   r   r   ÚScaleDegreesA  s   rt   r    r   é(   r;   r   r   N)r   r   rX   r	   rI   r<   rJ   ©rt   rZ   r[   r\   r]   r   r   r   Ú test_compound_input_units_strict<  s     rw   c               	   C   s<  G dd„ dt ƒ} | dƒ}d|_t dƒ}||B }| d¡}|dƒ}t|dtj ƒ |dtj ƒ}t|dtj ƒ tjt	t
 ddd	¡d
�� |dtj ƒ}W 5 Q R X d|_||B }| d¡}tjt	t
 ddd	¡d
�� |dƒ}W 5 Q R X d|_||B }| d¡}|dƒ}t|dtj ƒ |dtj ƒ}t|dtj ƒ tjt	t
 ddd	¡d
�� |dtj ƒ}W 5 Q R X d|_||B }tjt	t
 ddd	¡d
�� |dƒ}W 5 Q R X d|_| dƒ}d|_| dƒ}d|_||@ }| d¡}|ddtj ƒ}t|d dtj ƒ t|d dtj ƒ tjt	t
 ddd	¡d
�� |ddƒ}W 5 Q R X dS )zL
    Test setting input_units_allow_dimensionless on one of the models.
    c                   @   s   e Zd ZdejiZdS )zCtest_compound_input_units_allow_dimensionless.<locals>.ScaleDegreesr$   Nrr   r   r   r   r   rt   `  s   rt   r    Trn   r   ru   r   r<   r   Fr   r   r;   r   N)r   rE   ro   r	   r   ÚonerX   r   r   r
   r   r   r   rv   r   r   r   Ú-test_compound_input_units_allow_dimensionless[  sT    



ry   c                  C   sB   G dd„ dt ƒ} tƒ | ƒ B }t|dtj dtj ƒdtj ƒ dS )zn
    Test that return_units on the first model in the chain is respected for the
    input to the second.
    c                       sD   e Zd ZdZdZ‡ fdd„Zedd„ ƒZedd„ ƒZdd	„ Z	‡  Z
S )
z-test_compound_return_units.<locals>.PassModelr    c                    s   t ƒ j||Ž d S r4   )ÚsuperÚ__init__)r,   ÚargsÚkwargs©Ú	__class__r   r   r{   «  s    z6test_compound_return_units.<locals>.PassModel.__init__c                 S   s   t jt jdœS )zInput units.rj   ©r   r<   r=   r   r   r   rs   ®  s    z9test_compound_return_units.<locals>.PassModel.input_unitsc                 S   s   t jt jdœS )zOutput units.rj   r€   r=   r   r   r   Úreturn_units³  s    z:test_compound_return_units.<locals>.PassModel.return_unitsc                 S   s   |j |j fS r4   )r   )r,   r$   Úyr   r   r   r-   ¸  s    z6test_compound_return_units.<locals>.PassModel.evaluate)r.   r/   r0   r1   r2   r{   Úpropertyrs   r�   r-   Ú__classcell__r   r   r~   r   Ú	PassModel§  s   

r…   r   )r   éZ   N)r   r   r	   r   r<   )r…   r\   r   r   r   Útest_compound_return_units¡  s    r‡   )$Ú__doc__ÚnumpyrG   r   Znumpy.testingr   Zastropyr   r   Zastropy.modeling.corer   Zastropy.modeling.modelsr   r   r   r   Zastropy.tests.helperr	   Zastropy.unitsr
   r   r   r'   r(   r3   r^   r`   ra   rd   re   ri   rm   rq   rw   ry   r‡   r   r   r   r   Ú<module>   s2   ÿ*
R!F