plot.py 20 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/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
15
#
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
17

Martin Reinecke's avatar
Martin Reinecke committed
18
19
import os

20
21
import numpy as np

Martin Reinecke's avatar
fix    
Martin Reinecke committed
22
23
24
25
26
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 .field import Field
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
    xc, yc = (xsize-1)*0.5, (ysize-1)*0.5
Martin Reinecke's avatar
Martin Reinecke committed
44
    u, v = np.meshgrid(np.arange(xsize), np.arange(ysize))
Martin Reinecke's avatar
Martin Reinecke committed
45
    u, v = 2*(u-xc)/(xc/1.02), (v-yc)/(yc/1.02)
Martin Reinecke's avatar
Martin Reinecke committed
46
47
48
49
50
51
52
53
54

    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
55

56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
def _rgb_data(spectral_cube):
    _xyz = np.array(
          [[0.000160, 0.000662, 0.002362, 0.007242, 0.019110,
            0.043400, 0.084736, 0.140638, 0.204492, 0.264737,
            0.314679, 0.357719, 0.383734, 0.386726, 0.370702,
            0.342957, 0.302273, 0.254085, 0.195618, 0.132349,
            0.080507, 0.041072, 0.016172, 0.005132, 0.003816,
            0.015444, 0.037465, 0.071358, 0.117749, 0.172953,
            0.236491, 0.304213, 0.376772, 0.451584, 0.529826,
            0.616053, 0.705224, 0.793832, 0.878655, 0.951162,
            1.014160, 1.074300, 1.118520, 1.134300, 1.123990,
            1.089100, 1.030480, 0.950740, 0.856297, 0.754930,
            0.647467, 0.535110, 0.431567, 0.343690, 0.268329,
            0.204300, 0.152568, 0.112210, 0.081261, 0.057930,
            0.040851, 0.028623, 0.019941, 0.013842, 0.009577,
            0.006605, 0.004553, 0.003145, 0.002175, 0.001506,
            0.001045, 0.000727, 0.000508, 0.000356, 0.000251,
            0.000178, 0.000126, 0.000090, 0.000065, 0.000046,
            0.000033],
           [0.000017, 0.000072, 0.000253, 0.000769, 0.002004,
            0.004509, 0.008756, 0.014456, 0.021391, 0.029497,
            0.038676, 0.049602, 0.062077, 0.074704, 0.089456,
            0.106256, 0.128201, 0.152761, 0.185190, 0.219940,
            0.253589, 0.297665, 0.339133, 0.395379, 0.460777,
            0.531360, 0.606741, 0.685660, 0.761757, 0.823330,
            0.875211, 0.923810, 0.961988, 0.982200, 0.991761,
            0.999110, 0.997340, 0.982380, 0.955552, 0.915175,
            0.868934, 0.825623, 0.777405, 0.720353, 0.658341,
            0.593878, 0.527963, 0.461834, 0.398057, 0.339554,
            0.283493, 0.228254, 0.179828, 0.140211, 0.107633,
            0.081187, 0.060281, 0.044096, 0.031800, 0.022602,
            0.015905, 0.011130, 0.007749, 0.005375, 0.003718,
            0.002565, 0.001768, 0.001222, 0.000846, 0.000586,
            0.000407, 0.000284, 0.000199, 0.000140, 0.000098,
            0.000070, 0.000050, 0.000036, 0.000025, 0.000018,
            0.000013],
           [0.000705, 0.002928, 0.010482, 0.032344, 0.086011,
            0.197120, 0.389366, 0.656760, 0.972542, 1.282500,
            1.553480, 1.798500, 1.967280, 2.027300, 1.994800,
            1.900700, 1.745370, 1.554900, 1.317560, 1.030200,
            0.772125, 0.570060, 0.415254, 0.302356, 0.218502,
            0.159249, 0.112044, 0.082248, 0.060709, 0.043050,
            0.030451, 0.020584, 0.013676, 0.007918, 0.003988,
            0.001091, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
            0.000000]])

    MATRIX_SRGB_D65 = np.array(
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
111
            [[3.2404542, -1.5371385, -0.4985314],
112
             [-0.9692660,  1.8760108,  0.0415560],
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
113
             [0.0556434, -0.2040259,  1.0572252]])
114
115
116
117
118
119

    def _gammacorr(inp):
        mask = np.zeros(inp.shape, dtype=np.float64)
        mask[inp <= 0.0031308] = 1.
        r1 = 12.92*inp
        a = 0.055
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
120
        r2 = (1 + a) * (np.maximum(inp, 0.0031308) ** (1/2.4)) - a
121
122
123
        return r1*mask + r2*(1.-mask)

    def lambda2xyz(lam):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
124
125
126
        lammin = 380.
        lammax = 780.
        lam = np.asarray(lam, dtype=np.float64)
127
128
129
130
131
132
        lam = np.clip(lam, lammin, lammax)

        idx = (lam-lammin)/(lammax-lammin)*(_xyz.shape[1]-1)
        ii = np.maximum(0, np.minimum(79, int(idx)))
        w1 = 1.-(idx-ii)
        w2 = 1.-w1
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
133
        c = w1*_xyz[:, ii] + w2*_xyz[:, ii+1]
134
135
136
137
138
139
140
        return c

    def getxyz(n):
        E0, E1 = 1./700., 1./400.
        E = E0 + np.arange(n)*(E1-E0)/(n-1)
        res = np.zeros((3, n), dtype=np.float64)
        for i in range(n):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
141
            res[:, i] = lambda2xyz(1./E[i])
142
143
        return res

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
144
145
146
147
148
149
150
151
    def to_logscale(arr, lo, hi):
        res = arr.clip(lo, hi)
        res = np.log(res/hi)
        tmp = np.log(hi/lo)
        res += tmp
        res /= tmp
        return res

Philipp Arras's avatar
Philipp Arras committed
152
    shp = spectral_cube.shape[:-1]+(3,)
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
153
    spectral_cube = spectral_cube.reshape((-1, spectral_cube.shape[-1]))
154
155
    xyz = getxyz(spectral_cube.shape[-1])
    xyz_data = np.tensordot(spectral_cube, xyz, axes=[-1, -1])
Martin Reinecke's avatar
Martin Reinecke committed
156
157
    xyz_data /= xyz_data.max()
    xyz_data = to_logscale(xyz_data, max(1e-3, xyz_data.min()), 1.)
158
    rgb_data = xyz_data.copy()
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
159
160
    for x in range(xyz_data.shape[0]):
        rgb_data[x] = _gammacorr(np.matmul(MATRIX_SRGB_D65, xyz_data[x]))
Martin Reinecke's avatar
Martin Reinecke committed
161
    rgb_data = rgb_data.clip(0., 1.)
Philipp Arras's avatar
Philipp Arras committed
162
    return rgb_data.reshape(shp)
163
164


Martin Reinecke's avatar
Martin Reinecke committed
165
166
def _find_closest(A, target):
    # A must be sorted
Martin Reinecke's avatar
Martin Reinecke committed
167
168
    idx = np.clip(A.searchsorted(target), 1, len(A)-1)
    idx -= target - A[idx-1] < A[idx] - target
Martin Reinecke's avatar
Martin Reinecke committed
169
170
    return idx

Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
171

Philipp Arras's avatar
Philipp Arras committed
172
def _makeplot(name, block=True, dpi=None):
173
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
Martin Reinecke committed
174
    if name is None:
175
176
177
        plt.show(block=block)
        if block:
            plt.close()
Martin Reinecke's avatar
Martin Reinecke committed
178
179
        return
    extension = os.path.splitext(name)[1]
180
    if extension in (".pdf", ".png", ".svg"):
Martin Reinecke's avatar
Martin Reinecke committed
181
        args = {}
Philipp Arras's avatar
Philipp Arras committed
182
183
184
        if dpi is not None:
            args['dpi'] = float(dpi)
        plt.savefig(name, **args)
Martin Reinecke's avatar
Martin Reinecke committed
185
186
187
188
        plt.close()
    else:
        raise ValueError("file format not understood")

Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
189

Martin Reinecke's avatar
Martin Reinecke committed
190
def _limit_xy(**kwargs):
Martin Reinecke's avatar
Martin Reinecke committed
191
    import matplotlib.pyplot as plt
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
192
    x1, x2, y1, y2 = plt.axis()
clienhar's avatar
clienhar committed
193
194
195
196
    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
197
198
    plt.axis((x1, x2, y1, y2))

Martin Reinecke's avatar
Martin Reinecke committed
199

Martin Reinecke's avatar
Martin Reinecke committed
200
201
202
203
204
205
206
207
208
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
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    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
255
256
257

    plt.register_cmap(cmap=LinearSegmentedColormap("Planck-like", planckcmap))
    plt.register_cmap(cmap=LinearSegmentedColormap("High Energy", he_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
258
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Map", fd_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
259
260
    plt.register_cmap(cmap=LinearSegmentedColormap("Faraday Uncertainty",
                                                   fdu_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
261
    plt.register_cmap(cmap=LinearSegmentedColormap("Plus Minus", pm_cmap))
Martin Reinecke's avatar
Martin Reinecke committed
262

Martin Reinecke's avatar
Martin Reinecke committed
263

264
def _plot1D(f, ax, **kwargs):
265
    import matplotlib.pyplot as plt
266

267
268
269
270
271
    for i, fld in enumerate(f):
        if not isinstance(fld, Field):
            raise TypeError("incorrect data type")
        if i == 0:
            dom = fld.domain
272
273
            if (len(dom) != 1):
                raise ValueError("input field must have exactly one domain")
274
275
276
        else:
            if fld.domain != dom:
                raise ValueError("domain mismatch")
277
    dom = dom[0]
Martin Reinecke's avatar
Martin Reinecke committed
278

clienhar's avatar
clienhar committed
279
    label = kwargs.pop("label", None)
280
    if not isinstance(label, list):
Martin Reinecke's avatar
Martin Reinecke committed
281
        label = [label] * len(f)
Martin Reinecke's avatar
Martin Reinecke committed
282

Martin Reinecke's avatar
Martin Reinecke committed
283
    linewidth = kwargs.pop("linewidth", 1.)
Philipp Arras's avatar
Philipp Arras committed
284
    if not isinstance(linewidth, list):
Martin Reinecke's avatar
Martin Reinecke committed
285
        linewidth = [linewidth] * len(f)
Philipp Arras's avatar
Philipp Arras committed
286

clienhar's avatar
clienhar committed
287
    alpha = kwargs.pop("alpha", None)
Philipp Arras's avatar
Philipp Arras committed
288
    if not isinstance(alpha, list):
Martin Reinecke's avatar
Martin Reinecke committed
289
        alpha = [alpha] * len(f)
Philipp Arras's avatar
Philipp Arras committed
290

clienhar's avatar
clienhar committed
291
292
293
    ax.set_title(kwargs.pop("title", ""))
    ax.set_xlabel(kwargs.pop("xlabel", ""))
    ax.set_ylabel(kwargs.pop("ylabel", ""))
294

Martin Reinecke's avatar
Martin Reinecke committed
295
    if isinstance(dom, RGSpace):
296
        plt.yscale(kwargs.pop("yscale", "linear"))
297
298
299
300
        npoints = dom.shape[0]
        dist = dom.distances[0]
        xcoord = np.arange(npoints, dtype=np.float64)*dist
        for i, fld in enumerate(f):
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
301
            ycoord = fld.val
302
303
304
305
306
307
            plt.plot(xcoord, ycoord, label=label[i],
                     linewidth=linewidth[i], alpha=alpha[i])
        _limit_xy(**kwargs)
        if label != ([None]*len(f)):
            plt.legend()
        return
Martin Reinecke's avatar
Martin Reinecke committed
308
    elif isinstance(dom, PowerSpace):
309
310
        plt.xscale(kwargs.pop("xscale", "log"))
        plt.yscale(kwargs.pop("yscale", "log"))
Philipp Arras's avatar
Philipp Arras committed
311
        xcoord = dom.k_lengths
Martin Reinecke's avatar
Martin Reinecke committed
312
        for i, fld in enumerate(f):
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
313
            ycoord = fld.val.copy()
314
            ycoord[0] = ycoord[1]
Martin Reinecke's avatar
Martin Reinecke committed
315
316
            plt.plot(xcoord, ycoord, label=label[i],
                     linewidth=linewidth[i], alpha=alpha[i])
Martin Reinecke's avatar
Martin Reinecke committed
317
        _limit_xy(**kwargs)
318
319
        if label != ([None]*len(f)):
            plt.legend()
Martin Reinecke's avatar
Martin Reinecke committed
320
        return
321
322
323
324
325
326
327
328
    raise ValueError("Field type not(yet) supported")


def _plot2D(f, ax, **kwargs):
    import matplotlib.pyplot as plt

    dom = f.domain

329
330
331
332
333
334
335
336
    if len(dom) > 2:
        raise ValueError("DomainTuple can have at most two entries.")

    # check for multifrequency plotting
    have_rgb = False
    if len(dom) == 2:
        if (not isinstance(dom[1], RGSpace)) or len(dom[1].shape) != 1:
            raise TypeError("need 1D RGSpace as second domain")
337
        if dom[1].shape[0] == 1:
Martin Reinecke's avatar
Martin Reinecke committed
338
339
            from .sugar import makeField
            f = makeField(f.domain[0], f.val[..., 0])
340
        else:
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
341
            rgb = _rgb_data(f.val)
342
            have_rgb = True
343
344
345

    foo = kwargs.pop("norm", None)
    norm = {} if foo is None else {'norm': foo}
Philipp Arras's avatar
Philipp Arras committed
346
347

    foo = kwargs.pop("aspect", None)
348
    aspect = {} if foo is None else {'aspect': foo}
349
350
351
352
353

    ax.set_title(kwargs.pop("title", ""))
    ax.set_xlabel(kwargs.pop("xlabel", ""))
    ax.set_ylabel(kwargs.pop("ylabel", ""))
    dom = dom[0]
354
355
    if not have_rgb:
        cmap = kwargs.pop("colormap", plt.rcParams['image.cmap'])
356
357
358
359

    if isinstance(dom, RGSpace):
        nx, ny = dom.shape
        dx, dy = dom.distances
360
361
362
363
364
365
        if have_rgb:
            im = ax.imshow(
                rgb, extent=[0, nx*dx, 0, ny*dy], origin="lower", **norm,
                **aspect)
        else:
            im = ax.imshow(
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
366
                f.val.T, extent=[0, nx*dx, 0, ny*dy],
367
368
369
                vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
                cmap=cmap, origin="lower", **norm, **aspect)
            plt.colorbar(im)
370
371
        _limit_xy(**kwargs)
        return
Martin Reinecke's avatar
Martin Reinecke committed
372
    elif isinstance(dom, (HPSpace, GLSpace)):
Martin Reinecke's avatar
Martin Reinecke committed
373
374
375
        import pyHealpix
        xsize = 800
        res, mask, theta, phi = _mollweide_helper(xsize)
376
        if have_rgb:
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
377
378
            res = np.full(shape=res.shape+(3,), fill_value=1.,
                          dtype=np.float64)
379

Martin Reinecke's avatar
Martin Reinecke committed
380
381
382
383
        if isinstance(dom, HPSpace):
            ptg = np.empty((phi.size, 2), dtype=np.float64)
            ptg[:, 0] = theta
            ptg[:, 1] = phi
384
            base = pyHealpix.Healpix_Base(int(np.sqrt(dom.size//12)), "RING")
385
386
387
            if have_rgb:
                res[mask] = rgb[base.ang2pix(ptg)]
            else:
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
388
                res[mask] = f.val[base.ang2pix(ptg)]
Martin Reinecke's avatar
Martin Reinecke committed
389
390
391
392
393
394
        else:
            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)
395
396
397
            if have_rgb:
                res[mask] = rgb[ilat*dom[0].nlon + ilon]
            else:
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
398
                res[mask] = f.val[ilat*dom.nlon + ilon]
Martin Reinecke's avatar
Martin Reinecke committed
399
        plt.axis('off')
400
401
402
403
404
405
        if have_rgb:
            plt.imshow(res, origin="lower")
        else:
            plt.imshow(res, vmin=kwargs.get("zmin"), vmax=kwargs.get("zmax"),
                       cmap=cmap, origin="lower")
            plt.colorbar(orientation="horizontal")
406
407
408
409
410
411
412
413
414
415
416
417
        return
    raise ValueError("Field type not(yet) supported")


def _plot(f, ax, **kwargs):
    _register_cmaps()
    if isinstance(f, Field):
        f = [f]
    f = list(f)
    if len(f) == 0:
        raise ValueError("need something to plot")
    if not isinstance(f[0], Field):
Martin Reinecke's avatar
Martin Reinecke committed
418
        raise TypeError("incorrect data type")
419
    dom1 = f[0].domain
Martin Reinecke's avatar
Martin Reinecke committed
420
421
    if (len(dom1) == 1 and
        (isinstance(dom1[0], PowerSpace) or
422
423
         (isinstance(dom1[0], RGSpace) and
          len(dom1[0].shape) == 1))):
424
425
426
427
428
429
        _plot1D(f, ax, **kwargs)
        return
    else:
        if len(f) != 1:
            raise ValueError("need exactly one Field for 2D plot")
        _plot2D(f[0], ax, **kwargs)
Martin Reinecke's avatar
Martin Reinecke committed
430
431
        return
    raise ValueError("Field type not(yet) supported")
Martin Reinecke's avatar
Martin Reinecke committed
432

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
433

434
435
436
437
438
439
440
441
442
443
class Plot(object):
    def __init__(self):
        self._plots = []
        self._kwargs = []

    def add(self, f, **kwargs):
        """Add a figure to the current list of plots.

        Notes
        -----
Philipp Arras's avatar
Docs    
Philipp Arras committed
444
445
        After doing one or more calls `add()`, one needs to call `output()` to
        show or save the plot.
446
447
448

        Parameters
        ----------
Philipp Arras's avatar
Philipp Arras committed
449
        f: Field or list of Field
Philipp Arras's avatar
Philipp Arras committed
450
            If `f` is a single Field, it must be defined on a single `RGSpace`,
Martin Reinecke's avatar
typo    
Martin Reinecke committed
451
            `PowerSpace`, `HPSpace`, `GLSpace`.
Philipp Arras's avatar
Philipp Arras committed
452
            If it is a list, all list members must be Fields defined over the
453
454
            same one-dimensional `RGSpace` or `PowerSpace`.
        title: string
Philipp Arras's avatar
Docs    
Philipp Arras committed
455
            Title of the plot.
456
        xlabel: string
Philipp Arras's avatar
Philipp Arras committed
457
            Label for the x axis.
458
        ylabel: string
Philipp Arras's avatar
Philipp Arras committed
459
            Label for the y axis.
460
        [xyz]min, [xyz]max: float
Philipp Arras's avatar
Philipp Arras committed
461
            Limits for the values to plot.
462
        colormap: string
Philipp Arras's avatar
Philipp Arras committed
463
            Color map to use for the plot (if it is a 2D plot).
464
        linewidth: float or list of floats
Philipp Arras's avatar
Philipp Arras committed
465
            Line width.
466
        label: string of list of strings
Philipp Arras's avatar
Philipp Arras committed
467
            Annotation string.
468
        alpha: float or list of floats
Philipp Arras's avatar
Docs    
Philipp Arras committed
469
            Transparency value.
470
        """
Philipp Arras's avatar
Philipp Arras committed
471
472
        from .multi_field import MultiField
        if isinstance(f, MultiField):
Philipp Arras's avatar
Philipp Arras committed
473
474
475
476
477
478
479
480
481
            for kk in f.domain.keys():
                self._plots.append(f[kk])
                mykwargs = kwargs.copy()
                if 'title' in kwargs:
                    mykwargs['title'] = "{} {}".format(kk, kwargs['title'])
                else:
                    mykwargs['title'] = "{}".format(kk)
                self._kwargs.append(mykwargs)
            return
482
483
484
485
486
487
488
489
490
        self._plots.append(f)
        self._kwargs.append(kwargs)

    def output(self, **kwargs):
        """Plot the accumulated list of figures.

        Parameters
        ----------
        title: string
Philipp Arras's avatar
Philipp Arras committed
491
492
493
494
495
496
497
498
            Title of the full plot.
        nx, ny: int
            Number of subplots to use in x- and y-direction.
            Default: square root of the numer of plots, rounded up.
        xsize, ysize: float
            Dimensions of the full plot in inches. Default: 6.
        name: string
            If left empty, the plot will be shown on the screen,
499
            otherwise it will be written to a file with the given name.
Philipp Arras's avatar
Philipp Arras committed
500
            Supported extensions: .png and .pdf. Default: None.
501
502
503
        block: bool
            Override the blocking behavior of the non-interactive plotting
            mode. The plot will not be closed in this case but is left open!
504
505
506
507
508
509
        """
        import matplotlib.pyplot as plt
        nplot = len(self._plots)
        fig = plt.figure()
        if "title" in kwargs:
            plt.suptitle(kwargs.pop("title"))
510
511
512
513
514
515
516
517
        nx = kwargs.pop("nx", 0)
        ny = kwargs.pop("ny", 0)
        if nx == ny == 0:
            nx = ny = int(np.ceil(np.sqrt(nplot)))
        elif nx == 0:
            nx = np.ceil(nplot/ny)
        elif ny == 0:
            ny = np.ceil(nplot/nx)
518
519
520
521
522
523
524
525
526
527
528
529
        if nx*ny < nplot:
            raise ValueError(
                'Figure dimensions not sufficient for number of plots. '
                'Available plot slots: {}, number of plots: {}'
                .format(nx*ny, nplot))
        xsize = kwargs.pop("xsize", 6)
        ysize = kwargs.pop("ysize", 6)
        fig.set_size_inches(xsize, ysize)
        for i in range(nplot):
            ax = fig.add_subplot(ny, nx, i+1)
            _plot(self._plots[i], ax, **self._kwargs[i])
        fig.tight_layout()
Philipp Arras's avatar
Philipp Arras committed
530
        _makeplot(kwargs.pop("name", None), block=kwargs.pop("block", True), dpi=kwargs.pop("dpi", None))