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
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 *
Philipp Arras's avatar
Fixup  
Philipp Arras committed
26
from .. import Field, GLSpace, HPSpace, PowerSpace, RGSpace, dobj
27

Martin Reinecke's avatar
Martin Reinecke committed
28 29 30 31 32 33 34 35
# 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
36
# - labels
Martin Reinecke's avatar
Martin Reinecke committed
37

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
38

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

Martin Reinecke's avatar
Martin Reinecke committed
58 59 60 61 62 63 64 65 66
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
67

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

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
97

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

Martin Reinecke's avatar
Martin Reinecke committed
107

Martin Reinecke's avatar
Martin Reinecke committed
108 109 110 111 112 113 114 115 116
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
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 160 161 162
    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
163 164 165

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

Martin Reinecke's avatar
Martin Reinecke committed
171

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

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

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

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

211
    dom = dom[0]
Martin Reinecke's avatar
Martin Reinecke committed
212
    fig = plt.figure()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
213
    ax = fig.add_subplot(1, 1, 1)
Martin Reinecke's avatar
Martin Reinecke committed
214

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

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

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