There is a maintenance of MPCDF Gitlab on Thursday, April 22st 2020, 9:00 am CEST - Expect some service interruptions during this time

plot.py 11.8 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
from __future__ import absolute_import, division, print_function
20

Martin Reinecke's avatar
Martin Reinecke committed
21 22
import os

23 24 25
import numpy as np

from ..compat import *
Martin Reinecke's avatar
Martin Reinecke committed
26 27 28 29 30 31
from ..field import Field
from ..domains.gl_space import GLSpace
from ..domains.hp_space import HPSpace
from ..domains.power_space import PowerSpace
from ..domains.rg_space import RGSpace
from .. import dobj
32

Martin Reinecke's avatar
Martin Reinecke committed
33 34 35 36 37 38 39 40
# 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
41
# - labels
Martin Reinecke's avatar
Martin Reinecke committed
42

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
43

Martin Reinecke's avatar
Martin Reinecke committed
44 45 46
def _mollweide_helper(xsize):
    xsize = int(xsize)
    ysize = xsize//2
Martin Reinecke's avatar
Martin Reinecke committed
47
    res = np.full(shape=(ysize, xsize), fill_value=np.nan, dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
48 49 50 51 52 53 54 55 56 57 58 59 60 61
    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
62

Martin Reinecke's avatar
Martin Reinecke committed
63 64 65 66 67 68 69 70 71
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
72

Martin Reinecke's avatar
Martin Reinecke committed
73
def _makeplot(name):
74
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
75
    if dobj.rank != 0:
76
        plt.close()
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
77
        return
Martin Reinecke's avatar
Martin Reinecke committed
78 79
    if name is None:
        plt.show()
80
        plt.close()
Martin Reinecke's avatar
Martin Reinecke committed
81 82
        return
    extension = os.path.splitext(name)[1]
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
83
    if extension == ".pdf":
Martin Reinecke's avatar
Martin Reinecke committed
84 85
        plt.savefig(name)
        plt.close()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
86
    elif extension == ".png":
Martin Reinecke's avatar
Martin Reinecke committed
87 88
        plt.savefig(name)
        plt.close()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
89 90 91 92 93 94 95 96 97 98
    # 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
99 100 101
    else:
        raise ValueError("file format not understood")

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
102

Martin Reinecke's avatar
Martin Reinecke committed
103
def _limit_xy(**kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
104
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
105
    x1, x2, y1, y2 = plt.axis()
clienhar's avatar
clienhar committed
106 107 108 109
    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
110 111
    plt.axis((x1, x2, y1, y2))

Martin Reinecke's avatar
Martin Reinecke committed
112

Martin Reinecke's avatar
Martin Reinecke committed
113 114 115 116 117 118 119 120 121
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
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 160 161 162 163 164 165 166 167
    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
168 169 170

    plt.register_cmap(cmap=LinearSegmentedColormap("Planck-like", planckcmap))
    plt.register_cmap(cmap=LinearSegmentedColormap("High Energy", he_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
171
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Map", fd_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
172 173
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Uncertainty",
                                                   fdu_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
174
    plt.register_cmap(cmap=LinearSegmentedColormap("Plus Minus", pm_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
175

Martin Reinecke's avatar
Martin Reinecke committed
176

Martin Reinecke's avatar
Martin Reinecke committed
177
def _plot(f, ax, **kwargs):
178
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
179
    _register_cmaps()
180 181 182
    if isinstance(f, Field):
        f = [f]
    if not isinstance(f, list):
Martin Reinecke's avatar
Martin Reinecke committed
183
        raise TypeError("incorrect data type")
184 185 186 187 188 189 190 191 192 193 194
    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
195
                    (isinstance(dom[0], RGSpace) and len(dom[0].shape) == 1)):
196
                raise ValueError("PowerSpace or 1D RGSpace required")
Martin Reinecke's avatar
Martin Reinecke committed
197

clienhar's avatar
clienhar committed
198
    label = kwargs.pop("label", None)
Martin Reinecke's avatar
Martin Reinecke committed
199 200
    if label is None:
        label = [None] * len(f)
201
    if not isinstance(label, list):
Martin Reinecke's avatar
Martin Reinecke committed
202 203
        label = [label]

clienhar's avatar
clienhar committed
204
    linewidth = kwargs.pop("linewidth", None)
Philipp Arras's avatar
Philipp Arras committed
205
    if linewidth is None:
Martin Reinecke's avatar
Martin Reinecke committed
206
        linewidth = [1.] * len(f)
Philipp Arras's avatar
Philipp Arras committed
207 208 209
    if not isinstance(linewidth, list):
        linewidth = [linewidth]

clienhar's avatar
clienhar committed
210
    alpha = kwargs.pop("alpha", None)
Philipp Arras's avatar
Philipp Arras committed
211 212 213 214 215
    if alpha is None:
        alpha = [None] * len(f)
    if not isinstance(alpha, list):
        alpha = [alpha]

216
    dom = dom[0]
clienhar's avatar
clienhar committed
217 218 219 220
    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
221
    if isinstance(dom, RGSpace):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
222
        if len(dom.shape) == 1:
Martin Reinecke's avatar
Martin Reinecke committed
223 224
            npoints = dom.shape[0]
            dist = dom.distances[0]
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
225
            xcoord = np.arange(npoints, dtype=np.float64)*dist
Martin Reinecke's avatar
Martin Reinecke committed
226
            for i, fld in enumerate(f):
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
227
                ycoord = fld.to_global_data()
Martin Reinecke's avatar
Martin Reinecke committed
228 229
                plt.plot(xcoord, ycoord, label=label[i],
                         linewidth=linewidth[i], alpha=alpha[i])
Martin Reinecke's avatar
Martin Reinecke committed
230
            _limit_xy(**kwargs)
231 232
            if label != ([None]*len(f)):
                plt.legend()
Martin Reinecke's avatar
Martin Reinecke committed
233
            return
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
234
        elif len(dom.shape) == 2:
235
            f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
236 237 238 239
            nx = dom.shape[0]
            ny = dom.shape[1]
            dx = dom.distances[0]
            dy = dom.distances[1]
Philipp Arras's avatar
Philipp Arras committed
240 241
            xc = np.arange(nx, dtype=np.float64)*dx
            yc = np.arange(ny, dtype=np.float64)*dy
Martin Reinecke's avatar
tweaks  
Martin Reinecke committed
242
            im = ax.imshow(fld.to_global_data(),
Martin Reinecke's avatar
Martin Reinecke committed
243
                           extent=[xc[0], xc[-1], yc[0], yc[-1]],
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
244
                           vmin=kwargs.get("zmin"),
Martin Reinecke's avatar
Martin Reinecke committed
245
                           vmax=kwargs.get("zmax"), cmap=cmap, origin="lower")
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
246 247 248 249
            # 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
250
            plt.colorbar(im)
Martin Reinecke's avatar
Martin Reinecke committed
251
            _limit_xy(**kwargs)
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 266
        return
    elif isinstance(dom, HPSpace):
267
        f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
268 269 270 271 272 273 274
        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
275 276
        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
277
        plt.axis('off')
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
278
        plt.imshow(res, vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
Martin Reinecke's avatar
Martin Reinecke committed
279
                   cmap=cmap, origin="lower")
Martin Reinecke's avatar
Martin Reinecke committed
280
        plt.colorbar(orientation="horizontal")
Martin Reinecke's avatar
Martin Reinecke committed
281 282
        return
    elif isinstance(dom, GLSpace):
283
        f = f[0]
Martin Reinecke's avatar
Martin Reinecke committed
284 285 286 287 288 289 290 291
        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
292
        res[mask] = f.to_global_data()[ilat*dom.nlon + ilon]
Martin Reinecke's avatar
Martin Reinecke committed
293 294

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

    raise ValueError("Field type not(yet) supported")
Martin Reinecke's avatar
Martin Reinecke committed
301 302 303 304

_plots = []
_kwargs = []

Martin Reinecke's avatar
Martin Reinecke committed
305
def plot(f, **kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
306 307 308
    _plots.append(f)
    _kwargs.append(kwargs)

Martin Reinecke's avatar
Martin Reinecke committed
309 310
def plot_finish(nx, ny, **kwargs):
    global _plots, _kwargs
Martin Reinecke's avatar
Martin Reinecke committed
311 312 313
    import matplotlib.pyplot as plt
    nplot = len(_plots)
    fig = plt.figure()
Martin Reinecke's avatar
Martin Reinecke committed
314 315 316 317 318
    if "title" in kwargs:
        plt.suptitle(kwargs.pop("title"))
    xsize = kwargs.pop("xsize", 6)
    ysize = kwargs.pop("ysize", 6)
    fig.set_size_inches(xsize, ysize)
Martin Reinecke's avatar
Martin Reinecke committed
319
    for i in range(nplot):
Martin Reinecke's avatar
Martin Reinecke committed
320
        ax = fig.add_subplot(nx,ny,i+1)
Martin Reinecke's avatar
Martin Reinecke committed
321
        _plot(_plots[i], ax, **_kwargs[i])
Martin Reinecke's avatar
Martin Reinecke committed
322 323 324
    _makeplot(kwargs.pop("name", None))
    _plots = []
    _kwargs = []