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 *
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 89 90 91
        plt.savefig(name)
        plt.close()
    else:
        raise ValueError("file format not understood")

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
92

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

Martin Reinecke's avatar
Martin Reinecke committed
102

Martin Reinecke's avatar
Martin Reinecke committed
103 104 105 106 107 108 109 110 111
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
112 113 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
    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
158 159 160

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

Martin Reinecke's avatar
Martin Reinecke committed
166

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

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

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

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

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

        plt.axis('off')
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
285
        plt.imshow(res, vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
Martin Reinecke's avatar
Martin Reinecke committed
286
                   cmap=cmap, origin="lower")
Martin Reinecke's avatar
Martin Reinecke committed
287
        plt.colorbar(orientation="horizontal")
Martin Reinecke's avatar
Martin Reinecke committed
288 289 290
        return

    raise ValueError("Field type not(yet) supported")
Martin Reinecke's avatar
Martin Reinecke committed
291 292 293 294

_plots = []
_kwargs = []

Martin Reinecke's avatar
Martin Reinecke committed
295
def plot(f, **kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
296 297 298
    _plots.append(f)
    _kwargs.append(kwargs)

Martin Reinecke's avatar
Martin Reinecke committed
299
def plot_finish(**kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
300
    global _plots, _kwargs
Martin Reinecke's avatar
Martin Reinecke committed
301 302 303
    import matplotlib.pyplot as plt
    nplot = len(_plots)
    fig = plt.figure()
Martin Reinecke's avatar
Martin Reinecke committed
304 305
    if "title" in kwargs:
        plt.suptitle(kwargs.pop("title"))
Martin Reinecke's avatar
Martin Reinecke committed
306 307
    nx = kwargs.pop("nx", 1)
    ny = kwargs.pop("ny", 1)
Martin Reinecke's avatar
Martin Reinecke committed
308 309 310
    xsize = kwargs.pop("xsize", 6)
    ysize = kwargs.pop("ysize", 6)
    fig.set_size_inches(xsize, ysize)
Martin Reinecke's avatar
Martin Reinecke committed
311
    for i in range(nplot):
Martin Reinecke's avatar
Martin Reinecke committed
312
        ax = fig.add_subplot(ny,nx,i+1)
Martin Reinecke's avatar
Martin Reinecke committed
313
        _plot(_plots[i], ax, **_kwargs[i])
Martin Reinecke's avatar
Martin Reinecke committed
314 315 316
    _makeplot(kwargs.pop("name", None))
    _plots = []
    _kwargs = []