B
    zdư                 @   s  d dl Z e 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Zd dlZd dlmZ d dlZd dlZd dlZddlmZ ejZed	 ZG d
d dejjZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZ G dd deZ!G dd de!Z"G dd deZ#G dd de#Z$d4d$d%Z%d&d' Z&d(d) Z'd5d*d+Z(G d,d- d-ej)j*Z+G d.d/ d/ej,j-Z.G d0d1 d1ej,j-Z/G d2d3 d3ej,j-Z0dS )6    NZhealpy)deprecated_renamed_argument   )	projector)rotator)	pixelfunc)UNSEENg     f@c            
       s   e Zd ZdZ fddZdd Zdd Zdd	 Zd
d Zdd Z	dde
dddddddf
ddZdd Zd(ddZdd Zd)ddZd*ddZd+ddZeddd d,d"d#Zd$d% Zeddd d-d&d'Z  ZS ).SphericalProjAxesa   Define a special Axes to take care of spherical projection.

    Parameters
    ----------
    projection : a SphericalProj class or a class derived from it.
        type of projection
    rot : list or string
        define rotation. See rotator.
    coord : list or string
        define coordinate system. See rotator.
    coordprec : number of digit after floating point for coordinates display.
    format : format string for value display.

    Notes
    -----
    Other keywords from Axes (see Axes).
    c       
         s(  t |tjstd|f |dd |dd |dd d|di | _|dd |d	d
 d|d< tt| j	|| | 
d | d | j \}}}}| || | || | jjtd dt dd\}}	dt|d
 |	d
   | _dt d | _d| _i | _d| jd< d| jd< d S )Nz>First argument must be a SphericalProj class (or derived from)rotcoordflipconv)r	   r
   r   Z	arrayinfoformatz%g	coordprec   equalZaspectoffFg       @g      ?T)directg      0@g?   localg      >@dpar)
issubclassPZSphericalProj	TypeErrorpopproj
setdefaultsuperr   __init__axisZset_autoscale_on
get_extentset_xlimset_ylimZang2xypidtornpsqrt_segment_threshold_segment_step_rad
_do_border_gratdef)
selfZ	ProjClassargskwdsxminxmaxyminymaxZdxZdy)	__class__ \/work/yifan.wang/ringdown/master-ringdown-env/lib/python3.7/site-packages/healpy/projaxes.pyr   ;   s0    




zSphericalProjAxes.__init__c             C   s
   || _ |S )z'Set the format string for value display)_format)r)   fr1   r1   r2   
set_formatW   s    zSphericalProjAxes.set_formatc             C   s
   || _ dS )z@Set the number of digits after floating point for coord display.N)
_coordprec)r)   nr1   r1   r2   set_coordprec\   s    zSphericalProjAxes.set_coordprecc       
      C   s   | j d }| ||}|dks,t| r0dS tj|| jd\}}| ||}|dkrbd}d}nt|t	krrd}| j
j}|dkr|d ||||dd	 f }	n|d
 |||f }	|	S )z_Format the coordinate for display in status bar. Take projection
        into account.
        z at N )Zdecimalsz%sz%s @ z(%g, %g) in %sr      zlon=%g, lat=%g)r3   
get_lonlatr#   isnananyaroundr6   	get_valuetypestrr   Zcoordsysstr)
r)   xyr   poslonlatvalZcoordsysresr1   r1   r2   format_coord`   s     
zSphericalProjAxes.format_coordc             C   s   | j j||dd\}}||fS )zGGet the coordinate in the coord system of the image, in lon/lat in deg.T)lonlat)r   Zxy2ang)r)   rB   rC   rE   rF   r1   r1   r2   r;   v   s    zSphericalProjAxes.get_lonlatc             C   s~   t |  dk rdS |  d }| }| j||\}}|dksJ|dkrNdS |jtjjk	rn|j||f rndS |||f S dS )z(Get the value of the map at position x,yr   Nr   )	lenZ
get_imagesZ	get_arrayr   Zxy2ijmaskr#   manomask)r)   rB   rC   ZimZarrijr1   r1   r2   r?   {   s    zSphericalProjAxes.get_valueNgraywhitec          	   K   s  | j j|||
|d}t|t|B tj||dB  }|dk	r| j j|||
|d}d||tj k< t|j	dkrz|}nt
 |}y|dkr||  }W n tk
r   d}Y nX y|dkr||  }W n tk
r   d}Y nX ||kr|}||kr|d8 }|d7 }t|||| ||	||d\}}| j  }tj||}| j|f|||d	d
|d|}| j  \}}}}| || | || |S )aG  Project a map on the SphericalProjAxes.

        Parameters
        ----------
        map : array-like
          The map to project.
        vec2pix_func : function
          The function describing the pixelisation.
        vmin, vmax : float, scalars
          min and max value to use instead of min max of the map
        badval : float
          The value of the bad pixels
        badcolor : str
          Color to use to plot bad values
        bgcolor : str
          Color to use for background
        cmap : a color map
          The colormap to use (see matplotlib.cm)
        rot : sequence
          In the form (lon, lat, psi) (unit: degree):the center of the map is
          at (lon, lat) and rotated by angle psi around that direction.
        coord : {'G', 'E', 'C', None}
          The coordinate system of the map ('G','E' or 'C'), rotate
          the map if different from the axes coord syst.
        alpha : array-like
          The alpha (transparency) map.

        Notes
        -----
        Other keywords are transmitted to :func:`matplotlib.Axes.imshow`
        )r	   r
   )badvalNr   r   g        g      ?)cmapnormbadcolorbgcolorZnearestlower)ZextentrU   rV   interpolationoriginalpha)r   projmapr#   r<   isinfr   Zmask_badinfuniquesizeplt	Normalizemin
ValueErrormaxget_color_tabler   rN   Zmasked_valuesZimshowr   r    )r)   mapvec2pix_funcvminvmaxrT   rW   rX   rU   rV   r	   r
   r\   r+   imgwZ	alpha_imgcmnnextZaximgr,   r-   r.   r/   r1   r1   r2   r]      sP    /$



zSphericalProjAxes.projmapc             O   sl  d}t |dk rtdt |dkr8t|d \}}nt |dkrt|d tkrp|d }t|d \}}qt|d t|d  }}nVt |dkrt|d tk	rtdqt|d t|d  }}|d }ntd|d	d}|dk	r*tjt	|dd
}|
d |d d |d< | j|ddddd }|dd}tj|||d}	tj||dd|	}	| jj|	|ddd\}
}| j|
||d| jd\}
}g }xt|
|D ]\}}|dk	r<ytj|\}}}W n    tjj|\}}}Y nX |d| |d| |dk	r<|d| tjj||f|}| | || qW |S )a  projplot is a wrapper around :func:`matplotlib.Axes.plot` to take into account the
        spherical projection.

        You can call this function as::

           projplot(theta, phi)        # plot a line going through points at coord (theta, phi)
           projplot(theta, phi, 'bo')  # plot 'o' in blue at coord (theta, phi)
           projplot(thetaphi)          # plot a line going through points at coord (thetaphi[0], thetaphi[1])
           projplot(thetaphi, 'bx')    # idem but with blue 'x'

        Parameters
        ----------
        theta, phi : float, array-like
          Coordinates of point to plot. Can be put into one 2-d array, first line is
          then *theta* and second line is *phi*. See *lonlat* parameter for unit.
        fmt : str
          A format string (see :func:`matplotlib.Axes.plot` for details)
        lonlat : bool, optional
          If True, theta and phi are interpreted as longitude and latitude
          in degree, otherwise, as colatitude and longitude in radian
        coord : {'E', 'G', 'C', None}
          The coordinate system of the points, only used if the coordinate
          coordinate system of the Axes has been defined and in this
          case, a rotation is performed
        rot : None or sequence
          rotation to be applied =(lon, lat, psi) : lon, lat will be position of the
          new Z axis, and psi is rotation around this axis, all in degree.
          if None, no rotation is performed
        direct : bool
          if True, the rotation to center the projection is not
          taken into account

        Notes
        -----
        Other keywords are passed to :func:`matplotlib.Axes.plot`.

        See Also
        --------
        projscatter, projtext
        Nr   zNo argument givenr   r   r:   zThird argument must be a stringzThree args maximumr	   )copyg     V@r
   rK   rJ   F)rJ   Y)r	   r
   	eulertyper   )r   	threshold)rt   	linestylemarkercolor)rL   re   r#   asarrayr@   rA   r   r   array
atleast_1dresizer   mkcoordRdir2vecRotatorIvec2xy_make_segmentr%   zip
matplotlibaxesZ_process_plot_formatZ_axesr   linesZLine2Dadd_lineappend)r)   r*   r+   fmtthetaphir	   r
   rJ   vecrB   rC   Zthelinesxxyyru   rv   rw   lr1   r1   r2   projplot   sV    ) 






zSphericalProjAxes.projplotc             O   s<  t | jd}|r |||| f}|dkr8t|\}}nt|t| }}|dd}|dk	rtjt|dd}|d |d d |d< | j	
|ddddd	 }|d
d}	tj|||	d}
tj||dd|
}
| j	j|
|ddd\}}| j||f||}|r8t | ds(g | _| j||f |S )a  Projscatter is a wrapper around :func:`matplotlib.Axes.scatter` to take into account the
        spherical projection.

        You can call this function as::

           projscatter(theta, phi)     # plot points at coord (theta, phi)
           projplot(thetaphi)          # plot points at coord (thetaphi[0], thetaphi[1])

        Parameters
        ----------
        theta, phi : float, array-like
          Coordinates of point to plot. Can be put into one 2-d array, first line is
          then *theta* and second line is *phi*. See *lonlat* parameter for unit.
        lonlat : bool, optional
          If True, theta and phi are interpreted as longitude and latitude
          in degree, otherwise, as colatitude and longitude in radian
        coord : {'E', 'G', 'C', None}, optional
          The coordinate system of the points, only used if the coordinate
          coordinate system of the axes has been defined and in this
          case, a rotation is performed
        rot : None or sequence, optional
          rotation to be applied =(lon, lat, psi) : lon, lat will be position of the
          new Z axis, and psi is rotation around this axis, all in degree.
          if None, no rotation is performed
        direct : bool, optional
          if True, the rotation to center the projection is not
          taken into account

        Notes
        -----
        Other keywords are passed to :func:`matplotlib.Axes.plot`.

        See Also
        --------
        projplot, projtext
        ZzoomtoolNr	   r   )rq   r:   g     V@r
   rK   rJ   F)rJ   rr   )r	   r
   rs   r   )r   _scatter_data)hasattrZfigurerq   r#   rx   r   ry   rz   r{   r   r|   r}   r~   r   r   r   Zscatterr   r   )r)   r   r   r*   r+   Zsave_input_dataZ
input_datar	   r
   rJ   r   rB   rC   sr1   r1   r2   projscatterC  s,    %
zSphericalProjAxes.projscatterc             K   s   |dkrt |\}}nt |t | }}|dd}|dk	rpt jt |dd}|d |d d |d< | j|ddddd }|d	d
}tj	|||d}tj
||dd|}| jj||dd
d\}	}
| j|	|
|f|S )aQ  Projtext is a wrapper around :func:`matplotlib.Axes.text` to take into account the
        spherical projection.

        Parameters
        ----------
        theta, phi : float, array-like
          Coordinates of point to plot. Can be put into one 2-d array, first line is
          then *theta* and second line is *phi*. See *lonlat* parameter for unit.
        text : str
          The text to be displayed.
        lonlat : bool, optional
          If True, theta and phi are interpreted as longitude and latitude
          in degree, otherwise, as colatitude and longitude in radian
        coord : {'E', 'G', 'C', None}, optional
          The coordinate system of the points, only used if the coordinate
          coordinate system of the axes has been defined and in this
          case, a rotation is performed
        rot : None or sequence, optional
          rotation to be applied =(lon, lat, psi) : lon, lat will be position of the
          new Z axis, and psi is rotation around this axis, all in degree.
          if None, no rotation is performed
        direct : bool, optional
          if True, the rotation to center the projection is not
          taken into account

        Notes
        -----
        Other keywords are passed to :func:`matplotlib.Axes.text`.

        See Also
        --------
        projplot, projscatter
        Nr	   r   )rq   r:   g     V@r
   rK   rJ   F)rJ   rr   )r	   r
   rs   r   )r   )r#   rx   r   ry   rz   r{   r   r|   r}   r~   r   r   r   text)r)   r   r   r   r+   r	   r
   rJ   r   rB   rC   r1   r1   r2   projtext  s    "
zSphericalProjAxes.projtextc       	      C   s  |d kr| j }t|t| }}tt|d| d t|d| d  }t||kd }g }g }t|dkrt||d  }t||d  }|| || nt|dkrr||d|d   ||d|d   xVtt|d D ]B}|||| ||d    |||| ||d     qW |||d d   |||d d   n|| || ||fS )Nr   r   r   rK   )	r%   r#   rz   r$   ZrollwhererL   r   range)	r)   rB   rC   rt   Zd2rm   r   r   rP   r1   r1   r2   r     s.    .
$

zSphericalProjAxes._make_segmentc             C   st   |dkr|dkr|\}}}n|dks,|dkr4t dt|}| j }td||d  }tt||d  }||fS )a  Get the min and max value of theta of the parallel to cover the
        field of view.

        Input:
          - the normalized vector of the direction of the center of the
            projection, in the reference frame of the graticule.
        Return:
          - vmin,vmax : between 0 and pi, vmin<vmax, the interval of theta
                        for the parallels crossing the field of view
        Nz.Both vy and vz must be given or both not giveng        g       @)re   r#   arccosr   get_fovrf   rd   r!   )r)   vxvyvzafovrj   rk   r1   r1   r2   get_parallel_interval  s    

z'SphericalProjAxes.get_parallel_intervalc             C   s   |dkr|dkr|\}}}n|dks,|dkr4t d| j }t|}||d krbtj tjfS t|t |d krtj tjfS t|}t||}||| d  ||| d  fS )ay  Get the min and max value of phi of the meridians to cover the field
        of view.

        Input:
          - the normalized vector of the direction of the center of the
            projection, in the reference frame of the graticule.
        Return:
          - vmin,vmax : the interval of phi for the
                        meridians crossing the field of view.
        Nz.Both vy and vz must be given or both not giveng       @)	re   r   r   r#   r   r!   abssinZarctan2)r)   r   r   r   r   thZsthZphi0r1   r1   r2   get_meridian_interval  s    


z'SphericalProjAxes.get_meridian_intervalverbosez1.15.0Tc          	   K   s  ||||f}|}|dkr"| j d }|dkr4| j d }|dkr@|}t|t }t|t }|st| j }	tj| jj|dd	|	}
nd}	d}
|
dd|
dd }}|
dd|
d	d }}|rtd
 |t  t }|rtd
 |t  t }|r|d d d t }|r*|d d d t }| |
\}}| |
\}}|rP|}|rZ|}|rd|}|rn|}td|t |t |t |t  |
dds| ||||||\}}tt||d|  || | }tt||d|  || | }t||t|| d | j}t||t|| d | j}d}g }|dd |dd x|D ]}t|td
  dk rd}d}n<t|dk rd}d}n$t|t dk rtd }d}nd}|| j|d | ||f||d| qdW |sD|td
 krDtd |krD|| j|d td
  |df||d| xP|D ]H}t|dk rbd}nd}|| j||d | |f||d| qJW t| drZ| jrZtdd t }|| j||d t d!ddd" || j||d d#t  d!ddd" td$dt }|| j|d d |d!ddd" || j|d t d |d!ddd" t| d%rz| j|||f n|||fg| _||fS )&a)  Draw a graticule.

        Input:
         - dpar: angular separation between parallels in degree
         - dmer: angular separation between meridians in degree
         - coord: coordinate system of the graticule ('G', 'E' or 'C')
         - local: if True, no rotation performed at all
        Nr   r   )r
   )r   r   r   pmaxpminmminmmaxg       @g     f@ih  r   z{0} {1} {2} {3}forceFg      ?g      Y@lwr   rw   kg|=-T:g        )r
   r   r   r'   r      z-k)r   r   gH.?iL_graticules)r(   r   r"   r}   r~   r   Z
get_centerr   r|   r   r   r!   r   r   loginfor   _get_interv_graticuler#   r>   arangerd   r&   r   r   r   r   r'   r   )r)   r   dmerr
   r   r   r+   ZgratargsZgratkwdsr   Zvec0Zu_pminZu_pmaxZu_mminZu_mmaxr   r   r   r   Z
theta_listZphi_listr   r   ZequatorZ	gratlinestr   pr1   r1   r2   	graticule  s    

"$$
("&
(""zSphericalProjAxes.graticulec             C   sf   t | drbxR| jD ]H\}}}x<|D ]4}x.|D ]&}|| jkrH| j| q,td q,W q"W qW | `dS )z5Delete all graticules previously created on the Axes.r   zline not in linesN)r   r   r   remover   warning)r)   Zdum1Zdum2gglr   r1   r1   r2   delgraticulesq  s    



zSphericalProjAxes.delgraticulesc             C   s  ddd}d}	d}
|| | }|| | }||	krL||| t  |	d t  }||
krr||| t  |
d ddt  }|| dk s|| d	krt|| }}ttt|t  d
}|t  | d }td|| ttt|t  d
}|t  | d }td|| ||fS )Nr   c             S   sx   d}| | dk r | d9 } d}d}| | }|| }t t |}t |d|  d|  | }|rtdt d|  }|S )NFg      ?<   Tr   
   g      N@)r#   floorlog10r>   )dr7   ro   ZarcminrB   rC   exzr1   r1   r2   set_prec  s    z9SphericalProjAxes._get_interv_graticule.<locals>.set_prec   $   r   )ro   g?g      @r   g      N@z/The interval between parallels is %d deg %.2f'.z/The interval between meridians is %d deg %.2f'.)r   )r"   rf   intr#   r   r>   r   r   )r)   r   r   r   r   r   r   r   r   Z	max_n_parZ	max_n_merZn_parZn_merZvdegZvarcminr1   r1   r2   r   }  s$    
z'SphericalProjAxes._get_interv_graticule)N)N)NN)NN)NNNNT)T)__name__
__module____qualname____doc__r   r5   r8   rI   r;   r?   r   r]   r   r   r   r   r   r   r   r   r   r   __classcell__r1   r1   )r0   r2   r   (   s8   P]
=2



t
r   c                   s.   e Zd ZdZ fddZd	 fdd	Z  ZS )
GnomonicAxesaC  Define a gnomonic Axes to handle gnomonic projection.

    Input:
      - rot=, coord= : define rotation and coordinate system. See rotator.
      - coordprec= : number of digit after floating point for coordinates display.
      - format= : format string for value display.

      Other keywords from Axes (see Axes).
    c                sD   | dd tt| jtjf|| d| _d| jd< d| jd< d S )Nr   r:   FTr   g      ?r   )r   r   r   r   r   ZGnomonicProjr'   r(   )r)   r*   r+   )r0   r1   r2   r     s
    
zGnomonicAxes.__init__   N      ?c                s(   | j j|||d tt| j||f|S )N)xsizeysizereso)r   set_proj_plane_infor   r   r]   )r)   rh   ri   r   r   r   r+   )r0   r1   r2   r]     s    zGnomonicAxes.projmap)r   Nr   )r   r   r   r   r   r]   r   r1   r1   )r0   r2   r     s   	r   c                   s   e Zd Zd fdd	Z  ZS )HpxGnomonicAxesFc                sd   t t | fdd}|dd}|dd }|dd}tt| j||f|||d|S )	Nc                s   t j| || dS )N)nest)r   vec2pix)rB   rC   r   )r   nsider1   r2   <lambda>      z)HpxGnomonicAxes.projmap.<locals>.<lambda>r   r   r   r   g      ?)r   r   r   )r   
npix2nsideget_map_sizer   r   r   r]   )r)   rh   r   r+   r4   r   r   r   )r0   )r   r   r2   r]     s    
zHpxGnomonicAxes.projmap)F)r   r   r   r]   r   r1   r1   )r0   r2   r     s   r   c                   s.   e Zd ZdZ fddZd fdd	Z  ZS )MollweideAxesaE  Define a mollweide Axes to handle mollweide projection.

    Input:
      - rot=, coord= : define rotation and coordinate system. See rotator.
      - coordprec= : number of digit after floating point for coordinates display.
      - format= : format string for value display.

      Other keywords from Axes (see Axes).
    c                sB   | dd tt| jtjf|| | dd | dd d S )Nr   r   gGz gGz @g)\(g)\(?)r   r   r   r   r   ZMollweideProjr   r    )r)   r*   r+   )r0   r1   r2   r     s    zMollweideAxes.__init__   c                s@   | j j|d tt| j||f|}| dd | dd |S )N)r   gGz gGz @g)\(g)\(?)r   r   r   r   r]   r   r    )r)   rh   ri   r   r+   rl   )r0   r1   r2   r]     s
    zMollweideAxes.projmap)r   )r   r   r   r   r   r]   r   r1   r1   )r0   r2   r     s   	r   c                   s   e Zd Zd fdd	Z  ZS )HpxMollweideAxesFc                s4   t t | fdd}tt| j||f|S )Nc                s   t j| || dS )N)r   )r   r   )rB   rC   r   )r   r   r1   r2   r     r   z*HpxMollweideAxes.projmap.<locals>.<lambda>)r   r   r   r   r   r]   )r)   rh   r   r+   r4   )r0   )r   r   r2   r]     s    zHpxMollweideAxes.projmap)F)r   r   r   r]   r   r1   r1   )r0   r2   r     s   r   c                   s.   e Zd ZdZ fddZd fdd	Z  ZS )	CartesianAxesz;Define a cylindrical Axes to handle cylindrical projection.c                sD   | dd tt| jtjf|| d| _dt d | _d| _	d S )Nr   r   r   g?T)
r   r   r   r   r   ZCartesianProjr%   r!   r&   r'   )r)   r*   r+   )r0   r1   r2   r     s
    zCartesianAxes.__init__   Nc                s*   | j j||||d tt| j||f|S )N)r   r   lonralatra)r   r   r   r   r]   )r)   rh   ri   r   r   r   r   r+   )r0   r1   r2   r]     s    zCartesianAxes.projmap)r   NNN)r   r   r   r   r   r]   r   r1   r1   )r0   r2   r     s   r   c                   s   e Zd Zd fdd	Z  ZS )HpxCartesianAxesFc                s4   t t | fdd}tt| j||f|S )Nc                s   t j| || dS )N)r   )r   r   )rB   rC   r   )r   r   r1   r2   r     r   z*HpxCartesianAxes.projmap.<locals>.<lambda>)r   r   r   r   r   r]   )r)   rh   r   r+   r4   )r0   )r   r   r2   r]     s    zHpxCartesianAxes.projmap)F)r   r   r   r]   r   r1   r1   )r0   r2   r     s   r   c                   s.   e Zd ZdZ fddZd fdd	Z  ZS )	OrthographicAxesaB  Define an orthographic Axes to handle orthographic projection.

    Input:
    - rot=, coord= : define rotation and coordinate system. See rotator.
    - coordprec= : num of digits after floating point for coordinates display.
    - format= : format string for value display.

    Other keywords from Axes (see Axes).
    c                s6   | dd tt| jtjf|| d| _d| _d S )Nr   r   g{Gz?F)r   r   r   r   r   ZOrthographicProjr%   r'   )r)   r*   r+   )r0   r1   r2   r     s    zOrthographicAxes.__init__   Fc                sR   | j j||d tt| j||f|}|r0d}nd}| | | | dd |S )N)r   half_skyg)\(?gGz @g)\()r   r   r   r   r]   r   r    )r)   rh   ri   r   r   r+   rl   ratio)r0   r1   r2   r]     s    zOrthographicAxes.projmap)r   F)r   r   r   r   r   r]   r   r1   r1   )r0   r2   r     s   	r   c                   s   e Zd Zd fdd	Z  ZS )HpxOrthographicAxesFc                s2   t t| fdd}tt| j||f|S )Nc                s   t j| || dS )N)r   )r   r   )rB   rC   r   )r   r   r1   r2   r     r   z-HpxOrthographicAxes.projmap.<locals>.<lambda>)r   r   rL   r   r   r]   )r)   rh   r   r+   r4   )r0   )r   r   r2   r]     s    zHpxOrthographicAxes.projmap)F)r   r   r   r]   r   r1   r1   )r0   r2   r     s   r   c                   s.   e Zd ZdZ fddZd fd	d
	Z  ZS )AzimuthalAxesaz  Define an Azimuthal Axes to handle azimuthal equidistant or
       Lambert azimuthal equal-area projections.

    Input:
      - rot=, coord= : define rotation and coordinate system. See rotator.
      - coordprec= : number of digit after floating point for coordinates display.
      - format= : format string for value display.

      Other keywords from Axes (see Axes).
    c                s0   | dd tt| jtjf|| d| _d S )Nr   r:   F)r   r   r   r   r   ZAzimuthalProjr'   )r)   r*   r+   )r0   r1   r2   r   -  s    zAzimuthalAxes.__init__r   N      ?TFc       	         s,   | j j|||||d tt| j||f|S )N)r   r   r   lambr   )r   r   r   r   r]   )	r)   rh   ri   r   r   r   r   r   r+   )r0   r1   r2   r]   2  s    zAzimuthalAxes.projmap)r   Nr   TF)r   r   r   r   r   r]   r   r1   r1   )r0   r2   r   !  s   
	    r   c                   s   e Zd Zd fdd	Z  ZS )HpxAzimuthalAxesFc       	         sr   t t | fdd}|dd}|dd }|dd}|dd	}tt| j||f||||d
|S )Nc                s   t j| || dS )N)r   )r   r   )rB   rC   r   )r   r   r1   r2   r   F  r   z*HpxAzimuthalAxes.projmap.<locals>.<lambda>r   i   r   r   g      ?r   T)r   r   r   r   )r   r   r   r   r   r   r]   )	r)   rh   r   r+   r4   r   r   r   r   )r0   )r   r   r2   r]   D  s    
zHpxAzimuthalAxes.projmap)F)r   r   r   r]   r   r1   r1   )r0   r2   r   C  s   r   r   皙?rR   rS   c
             C   s   t |||	d}
t|tkr| dr4tdd}nn| dr`|atjj	t
tf| |dd}nB| drtjjd|||d	}n| d
rtdd}nd }|d krtdd}| |_||_|| |
|fS )N)rW   rX   r   F)clipZsymlog2T)rj   rk   r   Zsymlog)r   	linthreshlinscalebasehist)create_colormapr@   rA   rY   
startswithLogNorm2
linthresh_r   colorsZFuncNormsymlog_forwardsymlog_backwardZ
SymLogNorm
HistEqNormLinNorm2rj   rk   autoscale_None)rj   rk   rG   rU   rV   r   r   r   rW   rX   Znewcmapr1   r1   r2   rg   g  s0    


rg   c          	   C   s(   | t  }td|td||    S )zC
    Alternative symmetric logarithmic function used in Planck
    g      ?g      @)r   r#   r   r$   )mrB   r1   r1   r2   r     s    r   c             C   s$   d|  }|d d | }t | }|S )Nr   r   r   )r   )rC   r   rB   r   r1   r1   r2   r     s    r   c             C   s   t | tkrn| dkr`tjtjtjtd}tj||  d}tj	
t|d | }qtj| }n0t | tj	jtj	j
gkr| }ntjtjd }t|drtj	d|j|j}n
t|}||d || || |S )	aA  Create a new colormap with specified bad/background colors.

    Parameters
    ----------
    cmap : string or matplotlib.colors.Colormap
        type of colormap to create
    badcolor : string
        color for bad pixels (passed to set_bad)
    bgcolor : string
        color for background (passed to set_under)
    )ZplanckZ
planck_logZwmapdataz	_cmap.datg     o@z
image.cmap_segmentdatanewcmg      ?)r@   rA   ospathjoindirnameabspath__file__r   r   ZListedColormapr#   Zloadtxtrn   Zget_cmapZLinearSegmentedColormapZrcParamsr   r   Nrq   Zset_overZ	set_underZset_bad)rU   rW   rX   ZdatapathZ	cmap_pathZcmap0r  r1   r1   r2   r     s&    



r   c               @   s&   e Zd Zd	ddZdd Zdd ZdS )
BoundaryLocatorr   Nc             C   s    |dk rt d|| _|| _d S )Nr   z%Number of locs must be greater than 1)re   NlocsrV   )r)   r  rV   r1   r1   r2   r     s    zBoundaryLocator.__init__c             C   s   t jdk r| j \}}n| j \}}| jdkrpt|t	| j
t|t|  | j
d   }d| }n"|t	| j
||  | j
d   }|S )Nz0.98r   g      ?r   )r   __version__ZviewInterval
get_boundsr   Zget_view_intervalrV   r#   r   r   r
  )r)   rj   rk   locsr1   r1   r2   __call__  s    

,
"zBoundaryLocator.__call__c             C   sH   |    | j \}}||k r(|| }}||kr@|d8 }|d7 }||fS )Nr   )Zverify_intervalsZdataIntervalr  )r)   rj   rk   r1   r1   r2   	autoscale  s    
zBoundaryLocator.autoscale)r   N)r   r   r   r   r  r  r1   r1   r1   r2   r	    s   
r	  c               @   sH   e Zd ZdddZdddZdd Zd	d
 Zdd Zdd Zdd Z	dS )r   NFc             C   s$   t jj| ||| d | _d | _d S )N)r   r   rc   r   xvalyval)r)   rj   rk   r   r1   r1   r2   r     s    zHistEqNorm.__init__c       	      C   s  |d kr| j }t|r2d}tj|tj}nd}tj|gtj}| | t	| j
t	| j }}||kr~tdnz||krd| S |rtj|}tjjt |||||d}tjjt|| j| jtj|d}tj |t|j< |dkr
|d }|S )Nry   scalarz/minvalue must be less than or equal to maxvalueg        )rM   r   )r   r#   iterablerN   rx   astypefloat64ry   r   floatrj   rk   re   getmaskfilledZinterpr  r  r_   r^   r   )	r)   valuer   vtyperG   rj   rk   rM   resultr1   r1   r2   r    s,    


  
zHistEqNorm.__call__c             C   s   |   stdt|r,d}tj|}nd}tj|g}tjj| || j| jtj	|d}tj
 |t|j< |dkr|d }|S )NzNot invertible until scaledry   r  )rM   r   )scaledre   r#   r  rN   ry   
_lininterpr  r  r  r_   r^   r   )r)   r  r  rG   r  r1   r1   r2   inverse  s    
 zHistEqNorm.inversec             C   sZ   d}| j d kr| | _ d}| jd kr4| | _d}|sL| jd ksL| jd krV| | d S )NFT)rj   rd   rk   rf   r  r  _set_xyvals)r)   rG   changedr1   r1   r2   r   ,  s    



zHistEqNorm.autoscale_Nonec             C   s"   |  | _| | _| | d S )N)rd   rj   rf   rk   r  )r)   rG   r1   r1   r2   r  7  s    

zHistEqNorm.autoscalec             C   sr  t j| }t |j}|jt jjk	r4||jB }|j|  }|jdk r~t j	ddgt j
d| _t j	| j| jgt j
d| _d S t|jd d}|dk r|j}y"t j||| j| jfdd\}}W n0 tk
r   t j||| j| jfd	\}}Y nX |j|jd kr|d d
 }|t j
t|  }t dg| dgg| _t | jg|d|d |d    | jgg| _d S )Nr:   r   r   )Zdtype   i  T)binsr   new)r"  r   rK   g        g      ?g      ?)r#   rN   rx   Zravelr^   r   rM   rO   ra   ry   r  r  rj   rk   r  rd   Z	histogramr   r  r  sumZconcatenateZcumsum)r)   rG   r   rm   Zdata2r"  r   r1   r1   r2   r  <  s.    

"zHistEqNorm._set_xyvalsc             C   sr  t |dr"d}t|tj}nd}t|gtj}||}|d }|d ||t|d k< |d ||dk< t|dk|t|k @ }|| }	||	d  ||	 ||	d   ||	 ||	d    || ||	d     }
t||	 ||	d   dk}||	| d  |
|< t|| |d k}|d |
|< t|| |d k}|d |
|< |
||< |dkrn|d }|S )N__len__ry   r  r   rK   r   )r   r#   rx   r  r  ZsearchsortedrL   r   )r)   rB   Xrr   Zxtyper   idxr   ZwokZiokZyywokrm   ZwlZwhr1   r1   r2   r  Z  s.    


>
zHistEqNorm._lininterp)NNF)N)
r   r   r   r   r  r  r   r  r  r  r1   r1   r1   r2   r     s   

r   c               @   s*   e Zd ZdZd	ddZdd Zdd ZdS )
r   zA
    Normalize a given value to the 0-1 range on a log scale
    Nc       	      C   s  |d kr| j }t|r2d}tj|tj}nd}tj|gtj}tjt	|j
|}| | t| jt| j }}||krtdn|dkrtdn||krt|dt| S |rtj|}tjjt |||||d}tj|t| t|t|  }d|j
|j
dk < d|j
|j
d	k< tj |t	|j
< |jtjjk	rtd
|jt	|j
< |dkr|d }|S )Nry   r  z/minvalue must be less than or equal to maxvaluer   zvalues must all be positiveg        )rM   g      ?r   F)r   r#   r  rN   rx   r  r  ry   masked_wherer^   r   r   r  rj   rk   re   r@   r  r  r   r_   rM   rO   )	r)   r  r   r  rG   rj   rk   rM   r  r1   r1   r2   r  {  s8    



 *
zLogNorm2.__call__c             C   s>   | j dks| jdkr:tjt|j|}tjj	
| | dS )z'autoscale only None-valued vmin or vmaxN)rj   rk   r#   rN   r(  r^   r   r   r   rc   r   )r)   ArG   r1   r1   r2   r     s    zLogNorm2.autoscale_Nonec             C   sj   |   stdt| jt| j }}t|rRtj|}|tj	|| | S |t
|| | S d S )NzNot invertible until scaled)r  re   r  rj   rk   r#   r  rN   rx   powerpow)r)   r  rj   rk   rG   r1   r1   r2   r    s    
zLogNorm2.inverse)N)r   r   r   r   r  r   r  r1   r1   r1   r2   r   v  s   
#r   c               @   s*   e Zd ZdZd	ddZdd Zdd ZdS )
r   zA
    Normalize a given value to the 0-1 range on a lin scale
    Nc       
      C   sV  |d kr| j }t|r2d}tj|tj}nd}tj|gtj}t|j	}tj
||}| | t| jt| j }}||krtdn||krt|dt| S |rtj|}tjjt |||||d}|| d||   }	d|	j	|	j	dk < d|	j	|	j	dk< tj |	|< |	jtjjk	r@d	|	j|< |dkrR|	d }	|	S )
Nry   r  z/minvalue must be less than or equal to maxvalueg        )rM   g      ?r   r   F)r   r#   r  rN   rx   r  r  ry   r^   r   r(  r   r  rj   rk   re   r@   r  r  r_   rM   rO   )
r)   r  r   r  rG   Zwinfrj   rk   rM   r  r1   r1   r2   r    s6    


 

zLinNorm2.__call__c             C   s>   | j dks| jdkr:tjt|j|}tjj	
| | dS )z'autoscale only None-valued vmin or vmaxN)rj   rk   r#   rN   r(  r^   r   r   r   rc   r   )r)   r)  rG   r1   r1   r2   r     s    zLinNorm2.autoscale_Nonec             C   s`   |   stdt| jt| j }}t|rLtj|}||| |  S ||| |  S d S )NzNot invertible until scaled)	r  re   r  rj   rk   r#   r  rN   rx   )r)   r  rj   rk   rG   r1   r1   r2   r    s    
zLinNorm2.inverse)N)r   r   r   r   r  r   r  r1   r1   r1   r2   r     s   
"r   )NNr   r   r   rR   rS   )rR   rS   )1logging	getLoggerr   Zastropy.utils.decoratorsr   r9   r   r   r   r}   r   r   Zmatplotlib.axesZmatplotlib.pyplotZpyplotrb   numpyr#   rq   r  Z_healpy_pixel_libr   r!   r"   r   ZAxesr   r   r   r   r   r   r   r   r   r   r   rg   r   r   r   ZtickerZLocatorr	  r   rc   r   r   r   r1   r1   r1   r2   <module>   sX   
    |"(      
%
*' @