_commands.py 20.6 KB
Newer Older
1
# This file is part of pymagglobal
2
#
3
# Copyright (C) 2020 Helmholtz Centre Potsdam
4
# GFZ German Research Centre for Geosciences, Potsdam, Germany
5
# (https://www.gfz-potsdam.de)
6
#
7
# pymagglobal is free software: you can redistribute it and/or modify
8
# it under the terms of the GNU General Public License as published by
9
10
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
11
#
12
# pymagglobal is distributed in the hope that it will be useful,
13
14
15
# 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.
16
#
17
18
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
19

20
''' This **private** module contains the commands that are called by the CLI.
Maximilian Schanner's avatar
Maximilian Schanner committed
21
22
23
24
'''
import numpy as np

from cartopy import crs as ccrs
25
from matplotlib import pyplot as plt
Maximilian Schanner's avatar
Maximilian Schanner committed
26

27
28
from pymagglobal import core
from pymagglobal import utils
Maximilian Schanner's avatar
Maximilian Schanner committed
29

30
31
32
33
34
35
36
37
38
from urllib.error import URLError
from warnings import warn


class TimeOutWarning(Warning):
    '''Just to give the warning a meaningful name.
    '''
    pass

Maximilian Schanner's avatar
Maximilian Schanner committed
39
40

def args2times(args):
41
42
    '''Read the range of the model and create a linear array to evaluate it.
    Make use of the command line arguments `--every` and `--res`.
Maximilian Schanner's avatar
Maximilian Schanner committed
43
44
45
46

    Parameters
    ----------
    args : object
47
48
49
50
51
52
        The SimpleNamespace object returned by ArgumentParser.parse_args().

    Returns
    -------
    times : ndarray
        An equally spaced array in years.
Maximilian Schanner's avatar
Maximilian Schanner committed
53
    '''
54
55
56
57
58
    if args.begin is not None:
        if args.longterm:
            t_min = lt2yr(args.begin)
        else:
            t_min = args.begin
59
60
61
        args.model.valid_epoch(t_min)
    else:
        t_min = args.model.t_min
62
63
64
65
66
67

    if args.end is not None:
        if args.longterm:
            t_max = lt2yr(args.end)
        else:
            t_max = args.end
68
69
70
        args.model.valid_epoch(t_max)
    else:
        t_max = args.model.t_max
Maximilian Schanner's avatar
Maximilian Schanner committed
71
72

    if args.every is not None:
73
        times = np.arange(t_min, t_max + args.every, args.every)
74
75
76
77
    elif args.res == 'knots':
        inds = np.nonzero(np.logical_and(t_min < args.model.knots,
                                         args.model.knots < t_max))
        times = args.model.knots[inds]
Maximilian Schanner's avatar
Maximilian Schanner committed
78
    else:
79
        times = np.linspace(t_min, t_max, int(args.res))
Maximilian Schanner's avatar
Maximilian Schanner committed
80
81
82
83

    return times


84
def yr2lt(times):
85
86
87
88
89
90
91
92
93
94
95
96
97
    '''Translate times given in years AD into kilo-years `before present
    <https://en.wikipedia.org/wiki/Before_Present>`_ often used for
    longterm models.

    Parameters
    ----------
    times: float or int
        Years AD.

    Returns
    -------
    float
       Kilo years before present (backward counting from 1/1/1950).
Stefan Mauerberger's avatar
Stefan Mauerberger committed
98
99
100
101
102
103
104

    Examples
    --------
    >>> yr2lt(1950)
    0.0
    >>> yr2lt(-50)
    2.0
105
    '''
106
107
108
109
    return -(times - 1950) / 1000


def lt2yr(times):
Stefan Mauerberger's avatar
Stefan Mauerberger committed
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
    '''Translate kilo-years `before present
    <https://en.wikipedia.org/wiki/Before_Present>`_ into years AD.

    Parameters
    ----------
    times: float or int
       Kilo-years before present (backward counting from 1/1/1950).

    Returns
    -------
    float or int
        Years AD.

    Examples
    --------
    >>> lt2yr(0)
    1950
    >>> lt2yr(2.)
    -50.0
    '''
130
131
132
    return -times*1000 + 1950


133
def local_curve(args, fig=None):
134
    '''Handle the local command and create a local curve at a given
Maximilian Schanner's avatar
Maximilian Schanner committed
135
136
137
138
139
    location.

    Parameters
    ----------
    args : object
Maximilian Schanner's avatar
Maximilian Schanner committed
140
        The SimpleNamespace object returned by ArgumentParser.parse_args().
141
142
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.
Maximilian Schanner's avatar
Maximilian Schanner committed
143
144
145
146
147
148

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
149
        local curve.
Maximilian Schanner's avatar
Maximilian Schanner committed
150
151
152
    '''
    # get an array of times
    times = args2times(args)
153
    # create a local curve using the core function, check is performed
154
    # in args2times
155
    curves = core.local_curve(times, (args.lat, args.lon), args.model,
156
                              field_type=args.type, check=False)
Maximilian Schanner's avatar
Maximilian Schanner committed
157
    # output formats for dif and nez components
158
159
160
161
    fmts = {'dif': ('%.2f', '%2.6f', '%2.6f', '%1.7e'),
            'nez': ('%.2f', '%1.7e', '%1.7e', '%1.7e')}

    # if the --longterm flag is set, translate the years accordingly
Maximilian Schanner's avatar
Maximilian Schanner committed
162

163
164
    if args.longterm:
        times = yr2lt(times)
Maximilian Schanner's avatar
Maximilian Schanner committed
165
166
167
168
169
170
171
172
173
    # if an output is specified, save the result to text
    if args.output is not None:
        np.savetxt(args.output,
                   np.array([times,
                             curves[0],
                             curves[1],
                             curves[2]]).T,
                   fmt=fmts[args.type],
                   delimiter=',',
174
                   header=f'Local curves at ({args.lat}°, {args.lon}°) '
175
                          f'for {args.model.name}\n'
Maximilian Schanner's avatar
Maximilian Schanner committed
176
                          f'{args.t_label},'
Maximilian Schanner's avatar
Maximilian Schanner committed
177
178
179
                          f'{utils._labels[args.type][0]},'
                          f'{utils._labels[args.type][1]},'
                          f'{utils._labels[args.type][2]}')
Maximilian Schanner's avatar
Maximilian Schanner committed
180

181
    # if the --no-show flag is not set, plot the local curve
182
    if not args.no_show or args.savefig is not None:
183
184
        if fig is None:
            fig = plt.figure(figsize=(10, 7))
185
        fig.suptitle(f'Local curves at ({args.lat}°, {args.lon}°) '
186
                     f'for {args.model.name}')
Maximilian Schanner's avatar
Maximilian Schanner committed
187
188
189
190
191
192
193
        axs = np.empty(3, dtype=object)
        axs[0] = fig.add_subplot(221)
        axs[1] = fig.add_subplot(222)
        axs[2] = fig.add_subplot(212)

        for it in range(3):
            axs[it].plot(times, curves[it])
Maximilian Schanner's avatar
Maximilian Schanner committed
194
195
            axs[it].set_title(utils._names[args.type][it])
            axs[it].set_ylabel(utils._nicelabel(utils._labels[args.type][it]))
196
197
198
199
            if args.longterm:
                axs[it].invert_xaxis()
                axs[it].set_xlabel(args.t_label)
            else:
Maximilian Schanner's avatar
Maximilian Schanner committed
200
                axs[it].set_xlabel(utils._nicelabel(args.t_label))
Maximilian Schanner's avatar
Maximilian Schanner committed
201
        fig.tight_layout(rect=[0, 0.03, 1, 0.95])
202
        return fig
Maximilian Schanner's avatar
Maximilian Schanner committed
203
204


205
def dipole_series(args, fig=None):
Stefan Mauerberger's avatar
Stefan Mauerberger committed
206
207
    '''Handle the dipole command and create a dipole-moment time series for
    the given model.
Maximilian Schanner's avatar
Maximilian Schanner committed
208
209
210
211

    Parameters
    ----------
    args : object
Maximilian Schanner's avatar
Maximilian Schanner committed
212
        The object returned by ArgumentParser.parse_args().
213
214
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.
Maximilian Schanner's avatar
Maximilian Schanner committed
215
216
217
218
219
220
221

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        dipole series.
Maximilian Schanner's avatar
Maximilian Schanner committed
222
223
224
    '''
    # get an array of times
    times = args2times(args)
225
226
227
    # create the dipole-moment time series for the given model, check is
    # performed in args2times
    dip_ser = core.dipole_series(times, args.model, check=False)
228
229
230
    # if the --longterm flag is set, translate the years accordingly
    if args.longterm:
        times = yr2lt(times)
Maximilian Schanner's avatar
Maximilian Schanner committed
231
232
233
234
235
236

    # if an output is specified, save the result to text
    if args.output is not None:
        np.savetxt(args.output,
                   np.array([times,
                             dip_ser]).T,
237
                   fmt=('%.2f', '%1.7e'),
Maximilian Schanner's avatar
Maximilian Schanner committed
238
                   delimiter=',',
239
                   header=f'Dipole moment series for '
240
                          f'{args.model.name}\n'
241
                          f'{args.t_label},m [m^2 A]')
Maximilian Schanner's avatar
Maximilian Schanner committed
242
243

    # if the --no-show flag is not set, plot the time series
244
    if not args.no_show or args.savefig is not None:
245
246
        if fig is None:
            fig = plt.figure(figsize=(10, 7))
Maximilian Schanner's avatar
Maximilian Schanner committed
247
248
        ax = fig.add_subplot(111)
        ax.plot(times, dip_ser)
249
        ax.set_title(f'Dipole moment series for {args.model.name}')
250
        ax.set_ylabel(r'$m$ [$10^{22}$ m$^2$ A]')
251
252
253
254
        if args.longterm:
            ax.invert_xaxis()
            ax.set_xlabel(args.t_label)
        else:
Maximilian Schanner's avatar
Maximilian Schanner committed
255
            ax.set_xlabel(utils._nicelabel(args.t_label))
Maximilian Schanner's avatar
Maximilian Schanner committed
256
        fig.tight_layout()
257
        return fig
Maximilian Schanner's avatar
Maximilian Schanner committed
258
259


260
def coefficients_series(args):
261
262
263
264
265
266
    '''Handle the coeff-series command and provide the model time series for
    given coefficients.

    Parameters
    ----------
    args : object
Maximilian Schanner's avatar
Maximilian Schanner committed
267
268
269
270
271
272
273
274
        the object returned by ArgumentParser.parse_args().

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        coefficient series.
275
276
277
    '''
    inds = args.model.valid_degrees_orders(args.degree, args.order)
    times = args2times(args)
278
279
    # evaluate the splines at the epoch, check is performed in args2times
    _, _, coeffs = core.coefficients(times, args.model, check=False)
280
281
282

    if args.longterm:
        times = yr2lt(times)
283
284
285
286
    # if an output is specified, save the result to text
    if args.output is not None:
        s = ''
        fmt = ['%.2f']
Maximilian Schanner's avatar
Maximilian Schanner committed
287
        for ell, m in zip(args.degree, args.order):
288
289
290
291
292
293
294
295
296
297
298
299
            s = ''.join([s, f'({int(ell)}, {m}),'])
            fmt.append('%1.7e')
        np.savetxt(args.output,
                   np.concatenate((np.atleast_2d(times).T,
                                   coeffs[:, inds]),
                                  axis=1),
                   fmt=fmt,
                   delimiter=',',
                   header=f'Coefficient series for {args.model.name} in nT. '
                          f'The tuples give the degree and order of the '
                          f'series.\ntime [{args.t_unit}],' + s[:-1])
    # if the --no-show flag is not set, plot the results
300
    if not args.no_show or args.savefig is not None:
301
302
303
        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111)
        ax.set_title(f'Coefficients series for {args.model.name}')
304
        ax.set_ylabel(r'$g_\ell^m$ [nT]')
305
306
        for ind, ell, m in zip(inds, args.degree, args.order):
            ax.plot(times, coeffs[:, ind], label=rf'$\ell={ell}, m={m}$')
307
308
309
310
311
        if args.longterm:
            ax.invert_xaxis()
            ax.set_xlabel(args.t_label)
        else:
            ax.set_xlabel(utils._nicelabel(args.t_label))
312
313
        ax.legend()
        fig.tight_layout()
314
        return fig
315
316


317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
def secular_variations_series(args):
    '''Handle the secular-variations-series command and provide the model time
    series for given degree and order.

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args().

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        coefficient series.
    '''
    inds = args.model.valid_degrees_orders(args.degree, args.order)
    times = args2times(args)
    # evaluate the splines at the epoch, check is performed in args2times
    _, _, svs = core.secular_variation(times, args.model, check=False)

    if args.longterm:
        times = yr2lt(times)
    # if an output is specified, save the result to text
    if args.output is not None:
        s = ''
        fmt = ['%.2f']
        for ell, m in zip(args.degree, args.order):
            s = ''.join([s, f'({int(ell)}, {m}),'])
            fmt.append('%1.7e')
        np.savetxt(args.output,
                   np.concatenate((np.atleast_2d(times).T,
                                   svs[:, inds]),
                                  axis=1),
                   fmt=fmt,
                   delimiter=',',
                   header=f'Secular variation series for {args.model.name} in '
                          'nT/yr. The tuples give the degree and order of the '
                          f'series.\ntime [{args.t_unit}],' + s[:-1])
    # if the --no-show flag is not set, plot the results
    if not args.no_show or args.savefig is not None:
        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111)
        ax.set_title(f'Secular variation series for {args.model.name}')
        ax.set_ylabel(r'$\dot{g}_\ell^m$ [nT/yr]')
        for ind, ell, m in zip(inds, args.degree, args.order):
            ax.plot(times, svs[:, ind], label=rf'$\ell={ell}, m={m}$')
        if args.longterm:
            ax.invert_xaxis()
            ax.set_xlabel(args.t_label)
        else:
            ax.set_xlabel(utils._nicelabel(args.t_label))
        ax.legend()
        fig.tight_layout()
        return fig


374
375
def coefficients_epoch(args):
    '''Handle the coeffs-epoch command and provide the model coefficients for
Stefan Mauerberger's avatar
Stefan Mauerberger committed
376
    a given epoch.
Maximilian Schanner's avatar
Maximilian Schanner committed
377
378
379
380

    Parameters
    ----------
    args : object
Maximilian Schanner's avatar
Maximilian Schanner committed
381
382
383
384
385
386
387
388
        the object returned by ArgumentParser.parse_args().

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        coefficients.
Maximilian Schanner's avatar
Maximilian Schanner committed
389
390
    '''
    # evaluate the splines at the epoch
391
    ls, ms, coeffs = core.coefficients(args.epoch, args.model)
Maximilian Schanner's avatar
Maximilian Schanner committed
392
393
394
395
    # if an output is specified, save the result to text
    if args.output is not None:
        np.savetxt(args.output,
                   np.array([ls, ms, coeffs]).T,
396
                   fmt=('%d', '%d', '%1.7e'),
Maximilian Schanner's avatar
Maximilian Schanner committed
397
                   delimiter=',',
398
                   header=f'Coefficients for {args.model.name} in nT at epoch '
399
                          f'{args.epoch} {args.t_unit}\n'
Maximilian Schanner's avatar
Maximilian Schanner committed
400
                          f'l,m,g_l^m')
401
402
403
404
    if not args.no_show or args.savefig is not None:
        return plot_coeffs(coeffs, ls, ms, args)


405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
def secular_variations_epoch(args):
    '''Handle the secular-variation-epoch command and provide the model
    secular variation for a given epoch.

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args().

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        coefficients.
    '''
    # evaluate the splines at the epoch
    ls, ms, svs = core.secular_variation(args.epoch, args.model)
    # if an output is specified, save the result to text
    if args.output is not None:
        np.savetxt(args.output,
                   np.array([ls, ms, svs]).T,
                   fmt=('%d', '%d', '%1.7e'),
                   delimiter=',',
                   header=f'Secular variations for {args.model.name} in nT/yr '
                          f'at epoch {args.epoch} {args.t_unit}\n'
                          f'l,m,g_l^m')
    if not args.no_show or args.savefig is not None:
        return plot_coeffs(svs, ls, ms, args, unit='nT/yr',
                           name=r'Secular variations $\dot{g}_\ell^m$')


def plot_coeffs(gs, ls, ms, args, unit='nT', name=r'Coefficients $g_\ell^m$'):
Maximilian Schanner's avatar
Maximilian Schanner committed
438
    '''Plot Gauss coefficients.
439
440
441
442

    Parameters
    ----------
    gs : ndarray
Maximilian Schanner's avatar
Maximilian Schanner committed
443
        Gauss-coefficients.
444
    ls : ndarray
Maximilian Schanner's avatar
Maximilian Schanner committed
445
        SH degrees.
446
    ms : ndarray
Maximilian Schanner's avatar
Maximilian Schanner committed
447
        SH orders.
448
449
450
451
    unit : str, optional
        The unit of the coeffs. Default is nT.
    name : str, optional
        Headline name of what is being plotted.
Maximilian Schanner's avatar
Maximilian Schanner committed
452
453
454
455
456
457

    Returns
    -------
    matplotlib.figure.Figure
        matplotlib.figure.Figure object, containing a plot of the coefficients.
    '''
458
459
460
    ms = np.array(ms)
    lmax = np.max(ls)
    fig, axs = plt.subplots(ncols=1, nrows=lmax, sharex=True,
461
                            figsize=(10, 10),
462
                            subplot_kw={'xlim': (-lmax-0.1, lmax+0.1)})
463
    fig.suptitle(f'{name} for {args.model.name} at epoch '
464
                 f'{args.epoch} {args.t_unit}\n'
465
466
467
                 f'Units are {unit}.')
    for ell, ax in enumerate(axs, start=1):
        mask = np.equal(ls, ell)
468
        ax.scatter(ms[mask], gs[mask], marker='x')
469
470
        ymax = 1.1*np.abs(gs[mask]).max()
        ax.set_ylim(-ymax, ymax)
471
        ax.set_ylabel(fr'$\ell = {ell}$')
472
        ax.hlines(0, *ax.get_xlim(), color='gray', lw=0.5)
473
474
475
476
477
    axs[-1].set_xticks(range(-lmax, lmax+1))
    axs[-1].set_xlabel('$m$')
    fig.tight_layout(h_pad=0, rect=[0, 0.03, 1, 0.92])
    fig.subplots_adjust(hspace=0)
    return fig
478
479


480
def maps(args, fig=None):
Stefan Mauerberger's avatar
Stefan Mauerberger committed
481
482
    '''Handle the map command and create field map of the model for a given
    epoch.
Maximilian Schanner's avatar
Maximilian Schanner committed
483
484
485
486

    Parameters
    ----------
    args : object
Maximilian Schanner's avatar
Maximilian Schanner committed
487
        the object returned by ArgumentParser.parse_args().
488
489
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.
Maximilian Schanner's avatar
Maximilian Schanner committed
490
491
492
493
494
495
496

    Returns
    -------
    matplotlib.figure.Figure
        If the --no-show flag is not set or a savefig output is specified,
        return a matplotlib.figure.Figure object, containing a plot of the
        field map.
Maximilian Schanner's avatar
Maximilian Schanner committed
497
    '''
498
    # set up the points as expected by dsh_basis
499
    z_at = utils.get_z_at(args.res, t=args.epoch)
Maximilian Schanner's avatar
Maximilian Schanner committed
500
501
    # BUGFIX: white edges in cartopy
    z_at[1] -= 180
Maximilian Schanner's avatar
Maximilian Schanner committed
502
    # use the core function to get the field
Maximilian Schanner's avatar
Maximilian Schanner committed
503
    field = core.field(z_at, args.model, field_type=args.type)
Maximilian Schanner's avatar
Maximilian Schanner committed
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    # output formats for dif and nez components
    fmts = {'dif': ('%.1f', '%.1f', '%2.6f', '%2.6f', '%1.7e'),
            'nez': ('%.1f', '%.1f', '%1.7e', '%1.7e', '%1.7e')}

    # if an output is specified, save the result to text
    if args.output is not None:
        np.savetxt(args.output,
                   np.array([z_at[0],
                             z_at[1],
                             field[0],
                             field[1],
                             field[2]]).T,
                   fmt=fmts[args.type],
                   delimiter=',',
518
                   header=f'Field map for {args.model.name} at epoch '
519
                          f'{args.epoch} {args.t_unit}\n'
Maximilian Schanner's avatar
Maximilian Schanner committed
520
                          f'lat,lon,'
Maximilian Schanner's avatar
Maximilian Schanner committed
521
522
523
                          f'{utils._labels[args.type][0]},'
                          f'{utils._labels[args.type][1]},'
                          f'{utils._labels[args.type][2]}')
Maximilian Schanner's avatar
Maximilian Schanner committed
524
525

    # if the --no-show flag is not set, plot the map using cartopy
526
    if not args.no_show or args.savefig is not None:
Maximilian Schanner's avatar
Maximilian Schanner committed
527
528
        cmaps = ['RdBu', 'RdBu', 'RdBu']
        units = [r'$\mu$T', r'$\mu$T', r'$\mu$T']
529
530
        levels = [10, 10, 10]
        extends = ['neither', 'neither', 'neither']
Maximilian Schanner's avatar
Maximilian Schanner committed
531
532
        if args.type == 'dif':
            field[2] /= 1000
533

534
            levels[0] = [-40, -32, -24, -16, -8, 0, 8, 16, 24, 32, 40]
535
536
            levels[1] = np.linspace(-90, 90, 11)
            extends[0] = 'both'
Maximilian Schanner's avatar
Maximilian Schanner committed
537
538
539
540
541
            cmaps[2] = 'Blues'
            units[0] = r'deg.'
            units[1] = r'deg.'
        else:
            field /= 1000
Maximilian Schanner's avatar
Maximilian Schanner committed
542
            vmax = np.ceil(np.max(np.abs(field), axis=1))
Maximilian Schanner's avatar
Maximilian Schanner committed
543

Maximilian Schanner's avatar
Maximilian Schanner committed
544
545
546
547
            for it in range(3):
                levels[it] = np.linspace(-vmax[it], vmax[it], 11)

        cbar_hght = 0.08
548
549
        if fig is None:
            proj = ccrs.Mollweide()
Maximilian Schanner's avatar
Maximilian Schanner committed
550
            fig, axs = plt.subplots(1, 3, figsize=(13, 3.4),
551
                                    subplot_kw={'projection': proj})
Maximilian Schanner's avatar
Maximilian Schanner committed
552
553
            fig.tight_layout()
            fig.subplots_adjust(top=0.8, bottom=0.25, wspace=0.1)
554
555
556
557
558
            colaxs = []
            for it in range(3):
                bnds = axs[it].get_position().bounds

                colaxs.append(fig.add_axes([bnds[0],
Maximilian Schanner's avatar
Maximilian Schanner committed
559
                                            bnds[1]-0.06-cbar_hght,
560
561
562
563
564
565
566
567
568
569
                                            bnds[2],
                                            cbar_hght]))
            cbar_orientation = 'horizontal'
        else:
            axs = fig.get_axes()
            proj = axs[0].projection

            colaxs = []
            for it in range(3):
                bnds = axs[it].get_position().bounds
570
                colaxs.append(fig.add_axes([bnds[0],
Maximilian Schanner's avatar
Maximilian Schanner committed
571
                                            bnds[1]-0.02-0.5*cbar_hght,
572
                                            bnds[2],
Maximilian Schanner's avatar
Maximilian Schanner committed
573
                                            0.5*cbar_hght]))
574
            cbar_orientation = 'horizontal'
575

Maximilian Schanner's avatar
Maximilian Schanner committed
576
577
578
579
        plt_lat, plt_lon, _ = proj.transform_points(ccrs.Geodetic(),
                                                    z_at[1],
                                                    90 - z_at[0]).T

580
        fig.suptitle(f'Field maps for {args.model.name} at epoch '
581
                     f'{args.epoch} {args.t_unit}')
Maximilian Schanner's avatar
Maximilian Schanner committed
582
583

        for it in range(3):
584
585
586
587
            mappable = axs[it].tricontourf(plt_lat, plt_lon, field[it],
                                           cmap=cmaps[it],
                                           levels=levels[it],
                                           extend=extends[it])
588
589
590
591
592
593
            try:
                axs[it].coastlines(alpha=0.8, lw=0.5)
            except URLError:
                warn('Server connection timed out. Coastlines could not be '
                     'loaded.',
                     category=TimeOutWarning)
Maximilian Schanner's avatar
Maximilian Schanner committed
594
            axs[it].set_global()
Maximilian Schanner's avatar
Maximilian Schanner committed
595
            axs[it].set_title(f'{utils._names[args.type][it]} [{units[it]}]')
596
            colaxs[it].tick_params(labelsize=12)
597
            fig.colorbar(mappable,
598
599
                         cax=colaxs[it],
                         orientation=cbar_orientation)
600
        return fig