commands.py 12.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# This module contains the commands that are used by the main module of
# pymagglobal.
#
# Copyright (C) 2020 Helmholtz Centre Potsdam GFZ,
#     German Research Centre for Geosciences, Potsdam, Germany
#
# Cite as:
# TODO
#
# This file is part of pymagglobal.
#
# pymagglobal 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.
#
# pymagglobal 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 pymagglobal. If not, see <https://www.gnu.org/licenses/>.

'''Just a bunch of functions ...

.. todo:: Doc-string missing in module.
Maximilian Schanner's avatar
Maximilian Schanner committed
28
29
30
31
32
33
34
35
'''
from warnings import warn

import numpy as np

from cartopy import crs as ccrs
from matplotlib import pyplot as plt, colors, cm

36
from pyfield import REARTH
Maximilian Schanner's avatar
Maximilian Schanner committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

from . import pymagglobal
from . import utils


class OutOfRangeWarning(Warning):
    ''' just to give the warning a meaningful name '''
    pass


def valid_epoch(args):
    ''' check whether the epoch given via the command line is in the range of
    the model

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
55
56
57
58
59

    Returns
    -------
    float
        the epoch in yrs
Maximilian Schanner's avatar
Maximilian Schanner committed
60
61
62
63
    '''
    # get the range of the model from the model file
    with open(args.input, 'r') as fh:
        t_min, t_max = map(float, fh.readline().split()[0:2])
64
65
66
67
68
69
    if args.longterm:
        epoch = lt2yr(args.epoch)
    else:
        epoch = args.epoch
    if epoch < t_min or t_max < epoch:
        warn(f'Epoch {epoch} is outside of the models range [{t_min}, '
Maximilian Schanner's avatar
Maximilian Schanner committed
70
71
             f'{t_max}]. The result will be extrapolated and not robust.',
             category=OutOfRangeWarning)
72
    return epoch
Maximilian Schanner's avatar
Maximilian Schanner committed
73
74
75
76
77
78
79
80
81
82
83
84
85


def args2times(args):
    ''' read the range of the model and create a linear array to evaluate it.
    make use of the command line arguments --every and --res

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
    '''
    with open(args.input, 'r') as fh:
        t_min, t_max = map(float, fh.readline().split()[0:2])
86
87
88
89
90
91
92
93
94
95
96
    if args.begin is not None:
        if args.longterm:
            t_min = lt2yr(args.begin)
        else:
            t_min = args.begin

    if args.end is not None:
        if args.longterm:
            t_max = lt2yr(args.end)
        else:
            t_max = args.end
Maximilian Schanner's avatar
Maximilian Schanner committed
97
98

    if args.every is not None:
99
        times = np.arange(t_min, t_max + args.every, args.every)
Maximilian Schanner's avatar
Maximilian Schanner committed
100
101
102
103
104
105
    else:
        times = np.linspace(t_min, t_max, args.res)

    return times


106
107
108
109
110
111
112
113
114
115
116
def yr2lt(times):
    ''' translate times given as years to the format / age often used
    for longterm models '''
    return -(times - 1950) / 1000


def lt2yr(times):
    ''' translate times given in the longterm format to years '''
    return -times*1000 + 1950


Maximilian Schanner's avatar
Maximilian Schanner committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def master_curve(args):
    ''' handle the master command and create a master curve at a given
    location.

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
    '''
    # get an array of times
    times = args2times(args)
    # get the splines from the input file
    splines = pymagglobal.file2splines(args.input)
    # create a master curve using the core function
    curves = pymagglobal.master_curve(times, (args.lat, args.lon), splines,
                                      ser_type=args.type)
    # output formats for dif and nez components
134
135
136
137
    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
138

139
140
    if args.longterm:
        times = yr2lt(times)
Maximilian Schanner's avatar
Maximilian Schanner committed
141
142
143
144
145
146
147
148
149
150
151
    # 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=',',
                   header=f'Master curves at ({args.lat}°, {args.lon}°) '
                          f'for the model {args.model}\n'
152
                          f'{args.t_label}'
Maximilian Schanner's avatar
Maximilian Schanner committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
                          f'{utils.labels[args.type][0]},'
                          f'{utils.labels[args.type][1]},'
                          f'{utils.labels[args.type][2]}')

    # if the --no-show flag is not set, plot the master curve
    if not args.no_show:
        fig = plt.figure(figsize=(10, 7))
        fig.suptitle(f'Master curves at ({args.lat}°, {args.lon}°) '
                     f'for the model {args.model}')
        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])
            axs[it].set_title(utils.names[args.type][it])
            axs[it].set_ylabel(utils.nicelabel(utils.labels[args.type][it]))
171
172
173
174
175
            if args.longterm:
                axs[it].invert_xaxis()
                axs[it].set_xlabel(args.t_label)
            else:
                axs[it].set_xlabel(utils.nicelabel(args.t_label))
Maximilian Schanner's avatar
Maximilian Schanner committed
176
        fig.tight_layout(rect=[0, 0.03, 1, 0.95])
Stefan Mauerberger's avatar
Stefan Mauerberger committed
177
        plt.draw()
Maximilian Schanner's avatar
Maximilian Schanner committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194


def dipole_series(args):
    ''' handle the dipole command and create a dipole-moment time series for
    the given model

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
    '''
    # get an array of times
    times = args2times(args)
    # get the splines from the input file
    splines = pymagglobal.file2splines(args.input)
    # create the dipole-moment time series for the given model
    dip_ser = pymagglobal.dipole_series(times, splines)
195
196
197
    # if the --longterm flag is set, translate the years accordingly
    if args.longterm:
        times = yr2lt(times)
Maximilian Schanner's avatar
Maximilian Schanner committed
198
199
200
201
202
203

    # 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,
204
                   fmt=('%.2f', '%1.7e'),
Maximilian Schanner's avatar
Maximilian Schanner committed
205
206
207
                   delimiter=',',
                   header=f'Dipole moment series for for the model '
                          f'{args.model}\n'
208
                          f'{args.t_label},m [m^2 A]')
Maximilian Schanner's avatar
Maximilian Schanner committed
209
210
211
212
213
214
215
216

    # if the --no-show flag is not set, plot the time series
    if not args.no_show:
        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111)
        ax.plot(times, dip_ser)
        ax.set_title(f'Dipole moment series for for the model {args.model}')
        ax.set_ylabel(r'$m$ [m$^2$ A]')
217
218
219
220
221
        if args.longterm:
            ax.invert_xaxis()
            ax.set_xlabel(args.t_label)
        else:
            ax.set_xlabel(utils.nicelabel(args.t_label))
Maximilian Schanner's avatar
Maximilian Schanner committed
222
        fig.tight_layout()
Stefan Mauerberger's avatar
Stefan Mauerberger committed
223
        plt.draw()
Maximilian Schanner's avatar
Maximilian Schanner committed
224
225
226
227
228
229
230
231
232
233
234
235


def coefficients(args):
    ''' handle the coefficients command and provide the model coefficients for
    a given epoch

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
    '''
    # check whether the epoch is in the range of the model
236
    epoch = valid_epoch(args)
Maximilian Schanner's avatar
Maximilian Schanner committed
237
238
239
    # get the splines from the input file
    splines = pymagglobal.file2splines(args.input)
    # evaluate the splines at the epoch
240
    ls, ms, coeffs = pymagglobal.coefficients(epoch, splines)
Maximilian Schanner's avatar
Maximilian Schanner committed
241
242
243
244
245
246
247
    # 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,
                   fmt=('%.d', '%d', '%1.7e'),
                   delimiter=',',
                   header=f'Coefficients for the model {args.model} at epoch '
248
                          f'{args.epoch} {args.t_unit}\n'
Maximilian Schanner's avatar
Maximilian Schanner committed
249
250
251
252
253
                          f'l,m,g_l^m')
    # if no output is given and the --no-show flag is not set, print the
    # coefficients to the command line
    elif not args.no_show:
        print(f'Coefficients for the model {args.model} at epoch '
254
              f'{args.epoch} {args.t_unit}')
Maximilian Schanner's avatar
Maximilian Schanner committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
        right = len(str(ms[-1]))
        print(f'l\t' + ' '*(right-1) + 'm\t g_l^m')
        for l, m, g in zip(ls, ms, coeffs):
            print(f'{l:d}\t{m:>{right}d}\t{g:>11.4e}')


def maps(args):
    ''' handle the map command and create field map of the model for a given
    epoch

    Parameters
    ----------
    args : object
        the object returned by ArgumentParser.parse_args()
    '''
    # check whether the epoch is in the range of the model
271
    epoch = valid_epoch(args)
Maximilian Schanner's avatar
Maximilian Schanner committed
272
273
274
275
276
277
278
279
    # set up a grid of equidistant points on the sphere
    tps = utils.equi_sph(args.res)
    # set upt the points as expected by pyfield
    z_at = np.empty((4, tps.shape[1]), order='F')
    z_at[0] = np.rad2deg(tps[0])
    # BUGFIX: white edges in cartopy
    z_at[1] = np.rad2deg(tps[1]) - 180
    z_at[2] = REARTH
280
    z_at[3] = epoch
Maximilian Schanner's avatar
Maximilian Schanner committed
281
282
283
    # get the splines from the input file
    splines = pymagglobal.file2splines(args.input)
    # use the core function to get the field
284
    field = pymagglobal.field(z_at, splines)
Maximilian Schanner's avatar
Maximilian Schanner committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
    # convert the field if necessary
    if args.type == 'dif':
        field = np.array(utils.nez2dif(*field))
    # 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=',',
                   header=f'Field map for the model {args.model} at epoch '
303
                          f'{args.epoch} {args.t_unit}\n'
Maximilian Schanner's avatar
Maximilian Schanner committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
                          f'lat,lon,'
                          f'{utils.labels[args.type][0]},'
                          f'{utils.labels[args.type][1]},'
                          f'{utils.labels[args.type][2]}')

    # if the --no-show flag is not set, plot the map using cartopy
    if not args.no_show:
        vmaxs = np.max(np.abs(field), axis=1)
        vmins = -vmaxs
        cmaps = ['RdBu', 'RdBu', 'RdBu']
        units = [r'$\mu$T', r'$\mu$T', r'$\mu$T']
        if args.type == 'dif':
            field[2] /= 1000
            vmaxs[2] /= 1000
            vmins[2] = 0
            cmaps[2] = 'Blues'
            units[0] = r'deg.'
            units[1] = r'deg.'
        else:
            vmaxs /= 1000
            vmins /= 1000
            field /= 1000

        cbar_hght = 0.06
        proj = ccrs.Mollweide()
        plt_lat, plt_lon, _ = proj.transform_points(ccrs.Geodetic(),
                                                    z_at[1],
                                                    90 - z_at[0]).T

        fig, axs = plt.subplots(1, 3, figsize=(13, 3),
                                subplot_kw={'projection': proj})
        fig.suptitle(f'Field maps for the model {args.model} at epoch '
336
                     f'{args.epoch} {args.t_unit}')
Maximilian Schanner's avatar
Maximilian Schanner committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355

        for it in range(3):
            axs[it].tripcolor(plt_lat, plt_lon, field[it],
                              vmin=vmins[it], vmax=vmaxs[it],
                              rasterized=True, cmap=cmaps[it])
            axs[it].coastlines(alpha=0.8, lw=0.5)
            axs[it].set_global()
            axs[it].set_title(f'{utils.names[args.type][it]} [{units[it]}]')
            bnds = axs[it].get_position().bounds
            # colorbar for the mean
            colax = fig.add_axes([bnds[0],
                                  bnds[1]-0.1-cbar_hght,
                                  bnds[2],
                                  cbar_hght])
            norm = colors.Normalize(vmin=vmins[it],
                                    vmax=vmaxs[it])
            fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmaps[it]),
                         cax=colax,
                         orientation='horizontal')
Stefan Mauerberger's avatar
Stefan Mauerberger committed
356
        plt.draw()