1__all__ = ["mollview", "projplot"]
2
3import numpy as np
4from .pixelfunc import ang2pix, npix2nside
5from .rotator import Rotator
6from matplotlib.projections.geo import GeoAxes
7
8###### WARNING #################
9# this module is work in progress, the aim is to reimplement the healpy
10# plot functions using the new features of matplotlib and remove most
11# of the custom projection code
12
13
14class ThetaFormatterShiftPi(GeoAxes.ThetaFormatter):
15    """Shifts labelling by pi
16
17    Shifts labelling from -180,180 to 0-360"""
18
19    def __call__(self, x, pos=None):
20        if x != 0:
21            x *= -1
22        if x < 0:
23            x += 2 * np.pi
24        return super(ThetaFormatterShiftPi, self).__call__(x, pos)
25
26
27def lonlat(theta, phi):
28    """Converts theta and phi to longitude and latitude
29
30    From colatitude to latitude and from astro longitude to geo longitude"""
31
32    longitude = -1 * np.asarray(phi)
33    latitude = np.pi / 2 - np.asarray(theta)
34    return longitude, latitude
35
36
37def mollview(
38    m=None,
39    rot=None,
40    coord=None,
41    unit="",
42    xsize=1000,
43    nest=False,
44    min=None,
45    max=None,
46    flip="astro",
47    format="%g",
48    cbar=True,
49    cmap=None,
50    norm=None,
51    graticule=False,
52    graticule_labels=False,
53    **kwargs
54):
55    """Plot a healpix map (given as an array) in Mollweide projection.
56
57    Parameters
58    ----------
59    map : float, array-like or None
60      An array containing the map, supports masked maps, see the `ma` function.
61      If None, will display a blank map, useful for overplotting.
62    rot : scalar or sequence, optional
63      Describe the rotation to apply.
64      In the form (lon, lat, psi) (unit: degrees) : the point at
65      longitude *lon* and latitude *lat* will be at the center. An additional rotation
66      of angle *psi* around this direction is applied.
67    coord : sequence of character, optional
68      Either one of 'G', 'E' or 'C' to describe the coordinate
69      system of the map, or a sequence of 2 of these to rotate
70      the map from the first to the second coordinate system.
71    unit : str, optional
72      A text describing the unit of the data. Default: ''
73    xsize : int, optional
74      The size of the image. Default: 800
75    nest : bool, optional
76      If True, ordering scheme is NESTED. Default: False (RING)
77    min : float, optional
78      The minimum range value
79    max : float, optional
80      The maximum range value
81    flip : {'astro', 'geo'}, optional
82      Defines the convention of projection : 'astro' (default, east towards left, west towards right)
83      or 'geo' (east towards roght, west towards left)
84    format : str, optional
85      The format of the scale label. Default: '%g'
86    cbar : bool, optional
87      Display the colorbar. Default: True
88    norm : {'hist', 'log', None}
89      Color normalization, hist= histogram equalized color mapping,
90      log= logarithmic color mapping, default: None (linear color mapping)
91    kwargs : keywords
92      any additional keyword is passed to pcolormesh
93    graticule : bool
94      add graticule
95    graticule_labels : bool
96      longitude and latitude labels
97    """
98
99    # not implemented features
100    if not (norm is None):
101        raise NotImplementedError()
102
103    # Create the figure
104    import matplotlib.pyplot as plt
105
106    width = 8.5
107    fig = plt.figure(figsize=(width, width * .63))
108    ax = fig.add_subplot(111, projection="mollweide")
109    # FIXME: make a more general axes creation that works also with subplots
110    # ax = plt.gcf().add_axes((.125, .1, .9, .9), projection="mollweide")
111
112    # remove white space around the image
113    plt.subplots_adjust(left=0.02, right=0.98, top=0.95, bottom=0.05)
114    if graticule and graticule_labels:
115        plt.subplots_adjust(left=0.04, right=0.98, top=0.95, bottom=0.05)
116
117    if not m is None:
118        # auto min and max
119        if min is None:
120            min = m.min()
121        if max is None:
122            max = m.max()
123
124    # allow callers to override the hold state by passing hold=True|False
125    washold = ax.ishold()
126    hold = kwargs.pop("hold", None)
127    if hold is not None:
128        ax.hold(hold)
129
130    try:
131        ysize = xsize / 2
132        theta = np.linspace(np.pi, 0, ysize)
133        phi = np.linspace(-np.pi, np.pi, xsize)
134
135        longitude = np.radians(np.linspace(-180, 180, xsize))
136        if flip == "astro":
137            longitude = longitude[::-1]
138        latitude = np.radians(np.linspace(-90, 90, ysize))
139        # project the map to a rectangular matrix xsize x ysize
140        PHI, THETA = np.meshgrid(phi, theta)
141        # coord or rotation
142        if coord or rot:
143            r = Rotator(coord=coord, rot=rot, inv=True)
144            THETA, PHI = r(THETA.flatten(), PHI.flatten())
145            THETA = THETA.reshape(ysize, xsize)
146            PHI = PHI.reshape(ysize, xsize)
147        nside = npix2nside(len(m))
148        if not m is None:
149            grid_pix = ang2pix(nside, THETA, PHI, nest=nest)
150            grid_map = m[grid_pix]
151
152            # plot
153            ret = plt.pcolormesh(
154                longitude,
155                latitude,
156                grid_map,
157                vmin=min,
158                vmax=max,
159                rasterized=True,
160                **kwargs
161            )
162
163        # graticule
164        plt.grid(graticule)
165        if graticule:
166            longitude_grid_spacing = 60  # deg
167            ax.set_longitude_grid(longitude_grid_spacing)
168            if width < 10:
169                ax.set_latitude_grid(45)
170                ax.set_longitude_grid_ends(90)
171
172        if graticule_labels:
173            ax.xaxis.set_major_formatter(ThetaFormatterShiftPi(longitude_grid_spacing))
174        else:
175            # remove longitude and latitude labels
176            ax.xaxis.set_ticklabels([])
177            ax.yaxis.set_ticklabels([])
178
179        # colorbar
180        if cbar and not m is None:
181            cb = fig.colorbar(
182                ret, orientation="horizontal", shrink=.4, pad=0.05, ticks=[min, max]
183            )
184            cb.ax.xaxis.set_label_text(unit)
185            cb.ax.xaxis.labelpad = -8
186            # workaround for issue with viewers, see colorbar docstring
187            cb.solids.set_edgecolor("face")
188
189        plt.draw()
190    finally:
191        ax.hold(washold)
192
193    return ret
194
195
196def projplot(theta, phi, fmt=None, **kwargs):
197    """projplot is a wrapper around :func:`matplotlib.Axes.plot` to take into account the
198    spherical projection.
199
200    You can call this function as::
201
202       projplot(theta, phi)        # plot a line going through points at coord (theta, phi)
203       projplot(theta, phi, 'bo')  # plot 'o' in blue at coord (theta, phi)
204
205    Parameters
206    ----------
207    theta, phi : float, array-like
208      Coordinates of point to plot in radians.
209    fmt : str
210      A format string (see :func:`matplotlib.Axes.plot` for details)
211
212    Notes
213    -----
214    Other keywords are passed to :func:`matplotlib.Axes.plot`.
215
216    See Also
217    --------
218    projscatter, projtext
219    """
220    import matplotlib.pyplot as plt
221
222    longitude, latitude = lonlat(theta, phi)
223    if fmt is None:
224        ret = plt.plot(longitude, latitude, **kwargs)
225    else:
226        ret = plt.plot(longitude, latitude, fmt, **kwargs)
227    return ret
228