plot.py 11.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
Martin Reinecke's avatar
Martin Reinecke committed
14
# Copyright(C) 2013-2018 Max-Planck-Society
15 16 17 18
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.

19 20
from __future__ import absolute_import, division, print_function
from ..compat import *
Martin Reinecke's avatar
Martin Reinecke committed
21
import numpy as np
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
22
from ..import Field, RGSpace, HPSpace, GLSpace, PowerSpace, dobj
Martin Reinecke's avatar
Martin Reinecke committed
23 24 25 26 27 28 29 30 31 32
import os

# relevant properties:
# - x/y size
# - x/y/z log
# - x/y/z min/max
# - colorbar/colormap
# - axis on/off
# - title
# - axis labels
Martin Reinecke's avatar
Martin Reinecke committed
33
# - labels
Martin Reinecke's avatar
Martin Reinecke committed
34

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
35

Martin Reinecke's avatar
Martin Reinecke committed
36 37 38
def _mollweide_helper(xsize):
    xsize = int(xsize)
    ysize = xsize//2
Martin Reinecke's avatar
Martin Reinecke committed
39
    res = np.full(shape=(ysize, xsize), fill_value=np.nan, dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
40 41 42 43 44 45 46 47 48 49 50 51 52 53
    xc = (xsize-1)*0.5
    yc = (ysize-1)*0.5
    u, v = np.meshgrid(np.arange(xsize), np.arange(ysize))
    u = 2*(u-xc)/(xc/1.02)
    v = (v-yc)/(yc/1.02)

    mask = np.where((u*u*0.25 + v*v) <= 1.)
    t1 = v[mask]
    theta = 0.5*np.pi-(
        np.arcsin(2/np.pi*(np.arcsin(t1) + t1*np.sqrt((1.-t1)*(1+t1)))))
    phi = -0.5*np.pi*u[mask]/np.maximum(np.sqrt((1-t1)*(1+t1)), 1e-6)
    phi = np.where(phi < 0, phi+2*np.pi, phi)
    return res, mask, theta, phi

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
54

Martin Reinecke's avatar
Martin Reinecke committed
55 56 57 58 59 60 61 62 63
def _find_closest(A, target):
    # A must be sorted
    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return idx

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
64

Martin Reinecke's avatar
Martin Reinecke committed
65
def _makeplot(name):
66
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
67
    if dobj.rank != 0:
68
        plt.close()
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
69
        return
Martin Reinecke's avatar
Martin Reinecke committed
70 71
    if name is None:
        plt.show()
72
        plt.close()
Martin Reinecke's avatar
Martin Reinecke committed
73 74
        return
    extension = os.path.splitext(name)[1]
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
75
    if extension == ".pdf":
Martin Reinecke's avatar
Martin Reinecke committed
76 77
        plt.savefig(name)
        plt.close()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
78
    elif extension == ".png":
Martin Reinecke's avatar
Martin Reinecke committed
79 80
        plt.savefig(name)
        plt.close()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
81 82 83 84 85 86 87 88 89 90
    # elif extension==".html":
        # import mpld3
        # mpld3.save_html(plt.gcf(),fileobj=name,no_extras=True)
        # import plotly.offline as py
        # import plotly.tools as tls
        # plotly_fig = tls.mpl_to_plotly(plt.gcf())
        # py.plot(plotly_fig,filename=name)
        # py.plot_mpl(plt.gcf(),filename=name)
        # import bokeh
        # bokeh.mpl.to_bokeh(plt.gcf())
Martin Reinecke's avatar
Martin Reinecke committed
91 92 93
    else:
        raise ValueError("file format not understood")

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
94

Martin Reinecke's avatar
Martin Reinecke committed
95
def _limit_xy(**kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
96
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
97
    x1, x2, y1, y2 = plt.axis()
clienhar's avatar
clienhar committed
98 99 100 101
    x1 = kwargs.pop("xmin", x1)
    x2 = kwargs.pop("xmax", x2)
    y1 = kwargs.pop("ymin", y1)
    y2 = kwargs.pop("ymax", y2)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
102 103
    plt.axis((x1, x2, y1, y2))

Martin Reinecke's avatar
Martin Reinecke committed
104

Martin Reinecke's avatar
Martin Reinecke committed
105 106 107 108 109 110 111 112 113
def _register_cmaps():
    try:
        if _register_cmaps._cmaps_registered:
            return
    except AttributeError:
        _register_cmaps._cmaps_registered = True

    from matplotlib.colors import LinearSegmentedColormap
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159
    planckcmap = {'red':   ((0., 0., 0.), (.4, 0., 0.), (.5, 1., 1.),
                            (.7, 1., 1.), (.8, .83, .83), (.9, .67, .67),
                            (1., .5, .5)),
                  'green': ((0., 0., 0.), (.2, 0., 0.), (.3, .3, .3),
                            (.4, .7, .7), (.5, 1., 1.), (.6, .7, .7),
                            (.7, .3, .3), (.8, 0., 0.), (1., 0., 0.)),
                  'blue':  ((0., .5, .5), (.1, .67, .67), (.2, .83, .83),
                            (.3, 1., 1.), (.5, 1., 1.), (.6, 0., 0.),
                            (1., 0., 0.))}
    he_cmap = {'red':   ((0., 0., 0.), (.167, 0., 0.), (.333, .5, .5),
                         (.5, 1., 1.), (1., 1., 1.)),
               'green': ((0., 0., 0.), (.5, 0., 0.), (.667, .5, .5),
                         (.833, 1., 1.), (1., 1., 1.)),
               'blue':  ((0., 0., 0.), (.167, 1., 1.), (.333, .5, .5),
                         (.5, 0., 0.), (1., 1., 1.))}
    fd_cmap = {'red':   ((0., .35, .35), (.1, .4, .4), (.2, .25, .25),
                         (.41, .47, .47), (.5, .8, .8), (.56, .96, .96),
                         (.59, 1., 1.), (.74, .8, .8), (.8, .8, .8),
                         (.9, .5, .5), (1., .4, .4)),
               'green': ((0., 0., 0.), (.2, 0., 0.), (.362, .88, .88),
                         (.5, 1., 1.), (.638, .88, .88), (.8, .25, .25),
                         (.9, .3, .3), (1., .2, .2)),
               'blue':  ((0., .35, .35), (.1, .4, .4), (.2, .8, .8),
                         (.26, .8, .8), (.41, 1., 1.), (.44, .96, .96),
                         (.5, .8, .8), (.59, .47, .47), (.8, 0., 0.),
                         (1., 0., 0.))}
    fdu_cmap = {'red':   ((0., 1., 1.), (0.1, .8, .8), (.2, .65, .65),
                          (.41, .6, .6), (.5, .7, .7), (.56, .96, .96),
                          (.59, 1., 1.), (.74, .8, .8), (.8, .8, .8),
                          (.9, .5, .5), (1., .4, .4)),
                'green': ((0., .9, .9), (.362, .95, .95), (.5, 1., 1.),
                          (.638, .88, .88), (.8, .25, .25), (.9, .3, .3),
                          (1., .2, .2)),
                'blue':  ((0., 1., 1.), (.1, .8, .8), (.2, 1., 1.),
                          (.41, 1., 1.), (.44, .96, .96), (.5, .7, .7),
                          (.59, .42, .42), (.8, 0., 0.), (1., 0., 0.))}
    pm_cmap = {'red':   ((0., 1., 1.), (.1, .96, .96), (.2, .84, .84),
                         (.3, .64, .64), (.4, .36, .36), (.5, 0., 0.),
                         (1., 0., 0.)),
               'green': ((0., .5, .5), (.1, .32, .32), (.2, .18, .18),
                         (.3, .8, .8),  (.4, .2, .2), (.5, 0., 0.),
                         (.6, .2, .2), (.7, .8, .8), (.8, .18, .18),
                         (.9, .32, .32), (1., .5, .5)),
               'blue':  ((0., 0., 0.), (.5, 0., 0.), (.6, .36, .36),
                         (.7, .64, .64), (.8, .84, .84), (.9, .96, .96),
                         (1., 1., 1.))}
Martin Reinecke's avatar
Martin Reinecke committed
160 161 162

    plt.register_cmap(cmap=LinearSegmentedColormap("Planck-like", planckcmap))
    plt.register_cmap(cmap=LinearSegmentedColormap("High Energy", he_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
163
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Map", fd_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
164 165
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Uncertainty",
                                                   fdu_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
166
    plt.register_cmap(cmap=LinearSegmentedColormap("Plus Minus", pm_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
167

Martin Reinecke's avatar
Martin Reinecke committed
168

Martin Reinecke's avatar
Martin Reinecke committed
169
def plot(f, **kwargs):
170
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
171
    _register_cmaps()
172 173 174
    if isinstance(f, Field):
        f = [f]
    if not isinstance(f, list):
Martin Reinecke's avatar
Martin Reinecke committed
175
        raise TypeError("incorrect data type")
176 177 178 179 180 181 182 183 184 185 186
    for i, fld in enumerate(f):
        if not isinstance(fld, Field):
            raise TypeError("incorrect data type")
        if i == 0:
            dom = fld.domain
            if len(dom) != 1:
                raise ValueError("input field must have exactly one domain")
        else:
            if fld.domain != dom:
                raise ValueError("domain mismatch")
            if not (isinstance(dom[0], PowerSpace) or
187
                    (isinstance(dom[0], RGSpace) and len(dom[0].shape) == 1)):
188
                raise ValueError("PowerSpace or 1D RGSpace required")
Martin Reinecke's avatar
Martin Reinecke committed
189

clienhar's avatar
clienhar committed
190
    label = kwargs.pop("label", None)
Martin Reinecke's avatar
Martin Reinecke committed
191 192
    if label is None:
        label = [None] * len(f)
193
    if not isinstance(label, list):
Martin Reinecke's avatar
Martin Reinecke committed
194 195
        label = [label]

clienhar's avatar
clienhar committed
196
    linewidth = kwargs.pop("linewidth", None)
Philipp Arras's avatar
Philipp Arras committed
197
    if linewidth is None:
Martin Reinecke's avatar
Martin Reinecke committed
198
        linewidth = [1.] * len(f)
Philipp Arras's avatar
Philipp Arras committed
199 200 201
    if not isinstance(linewidth, list):
        linewidth = [linewidth]

clienhar's avatar
clienhar committed
202
    alpha = kwargs.pop("alpha", None)
Philipp Arras's avatar
Philipp Arras committed
203 204 205 206 207
    if alpha is None:
        alpha = [None] * len(f)
    if not isinstance(alpha, list):
        alpha = [alpha]

208
    dom = dom[0]
Martin Reinecke's avatar
Martin Reinecke committed
209
    fig = plt.figure()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
210
    ax = fig.add_subplot(1, 1, 1)
Martin Reinecke's avatar
Martin Reinecke committed
211

clienhar's avatar
clienhar committed
212 213
    xsize = kwargs.pop("xsize", 6)
    ysize = kwargs.pop("ysize", 6)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
214
    fig.set_size_inches(xsize, ysize)
clienhar's avatar
clienhar committed
215 216 217 218
    ax.set_title(kwargs.pop("title", ""))
    ax.set_xlabel(kwargs.pop("xlabel", ""))
    ax.set_ylabel(kwargs.pop("ylabel", ""))
    cmap = kwargs.pop("colormap", plt.rcParams['image.cmap'])
Martin Reinecke's avatar
Martin Reinecke committed
219
    if isinstance(dom, RGSpace):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
220
        if len(dom.shape) == 1:
Martin Reinecke's avatar
Martin Reinecke committed
221 222
            npoints = dom.shape[0]
            dist = dom.distances[0]
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
223
            xcoord = np.arange(npoints, dtype=np.float64)*dist
Martin Reinecke's avatar
Martin Reinecke committed
224
            for i, fld in enumerate(f):
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
225
                ycoord = fld.to_global_data()
Martin Reinecke's avatar
Martin Reinecke committed
226 227
                plt.plot(xcoord, ycoord, label=label[i],
                         linewidth=linewidth[i], alpha=alpha[i])
Martin Reinecke's avatar
Martin Reinecke committed
228
            _limit_xy(**kwargs)
229 230
            if label != ([None]*len(f)):
                plt.legend()
Martin Reinecke's avatar
Martin Reinecke committed
231
            _makeplot(kwargs.get("name"))
Martin Reinecke's avatar
Martin Reinecke committed
232
            return
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
233
        elif len(dom.shape) == 2:
234
            f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
235 236 237 238
            nx = dom.shape[0]
            ny = dom.shape[1]
            dx = dom.distances[0]
            dy = dom.distances[1]
Philipp Arras's avatar
Philipp Arras committed
239 240
            xc = np.arange(nx, dtype=np.float64)*dx
            yc = np.arange(ny, dtype=np.float64)*dy
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
241
            im = ax.imshow(fld.to_global_data(),
Martin Reinecke's avatar
Martin Reinecke committed
242
                           extent=[xc[0], xc[-1], yc[0], yc[-1]],
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
243
                           vmin=kwargs.get("zmin"),
Martin Reinecke's avatar
Martin Reinecke committed
244
                           vmax=kwargs.get("zmax"), cmap=cmap, origin="lower")
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
245 246 247 248
            # from mpl_toolkits.axes_grid1 import make_axes_locatable
            # divider = make_axes_locatable(ax)
            # cax = divider.append_axes("right", size="5%", pad=0.05)
            # plt.colorbar(im,cax=cax)
Martin Reinecke's avatar
Martin Reinecke committed
249
            plt.colorbar(im)
Martin Reinecke's avatar
Martin Reinecke committed
250 251
            _limit_xy(**kwargs)
            _makeplot(kwargs.get("name"))
Martin Reinecke's avatar
Martin Reinecke committed
252 253 254 255 256
            return
    elif isinstance(dom, PowerSpace):
        plt.xscale('log')
        plt.yscale('log')
        plt.title('power')
Philipp Arras's avatar
Philipp Arras committed
257
        xcoord = dom.k_lengths
Martin Reinecke's avatar
Martin Reinecke committed
258
        for i, fld in enumerate(f):
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
259
            ycoord = fld.to_global_data()
Martin Reinecke's avatar
Martin Reinecke committed
260 261
            plt.plot(xcoord, ycoord, label=label[i],
                     linewidth=linewidth[i], alpha=alpha[i])
Martin Reinecke's avatar
Martin Reinecke committed
262
        _limit_xy(**kwargs)
263 264
        if label != ([None]*len(f)):
            plt.legend()
Martin Reinecke's avatar
Martin Reinecke committed
265
        _makeplot(kwargs.get("name"))
Martin Reinecke's avatar
Martin Reinecke committed
266 267
        return
    elif isinstance(dom, HPSpace):
268
        f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
269 270 271 272 273 274 275
        import pyHealpix
        xsize = 800
        res, mask, theta, phi = _mollweide_helper(xsize)

        ptg = np.empty((phi.size, 2), dtype=np.float64)
        ptg[:, 0] = theta
        ptg[:, 1] = phi
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
276 277
        base = pyHealpix.Healpix_Base(int(np.sqrt(f.size//12)), "RING")
        res[mask] = f.to_global_data()[base.ang2pix(ptg)]
Martin Reinecke's avatar
Martin Reinecke committed
278
        plt.axis('off')
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
279
        plt.imshow(res, vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
Martin Reinecke's avatar
Martin Reinecke committed
280
                   cmap=cmap, origin="lower")
Martin Reinecke's avatar
Martin Reinecke committed
281
        plt.colorbar(orientation="horizontal")
Martin Reinecke's avatar
Martin Reinecke committed
282
        _makeplot(kwargs.get("name"))
Martin Reinecke's avatar
Martin Reinecke committed
283 284
        return
    elif isinstance(dom, GLSpace):
285
        f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
286 287 288 289 290 291 292 293
        import pyHealpix
        xsize = 800
        res, mask, theta, phi = _mollweide_helper(xsize)
        ra = np.linspace(0, 2*np.pi, dom.nlon+1)
        dec = pyHealpix.GL_thetas(dom.nlat)
        ilat = _find_closest(dec, theta)
        ilon = _find_closest(ra, phi)
        ilon = np.where(ilon == dom.nlon, 0, ilon)
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
294
        res[mask] = f.to_global_data()[ilat*dom.nlon + ilon]
Martin Reinecke's avatar
Martin Reinecke committed
295 296

        plt.axis('off')
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
297
        plt.imshow(res, vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
Martin Reinecke's avatar
Martin Reinecke committed
298
                   cmap=cmap, origin="lower")
Martin Reinecke's avatar
Martin Reinecke committed
299
        plt.colorbar(orientation="horizontal")
Martin Reinecke's avatar
Martin Reinecke committed
300
        _makeplot(kwargs.get("name"))
Martin Reinecke's avatar
Martin Reinecke committed
301 302 303
        return

    raise ValueError("Field type not(yet) supported")