_commands.py

# This file is part of pymagglobal
#
# Copyright (C) 2020 Helmholtz Centre Potsdam
# GFZ German Research Centre for Geosciences, Potsdam, Germany
# (https://www.gfz-potsdam.de)
#
# 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 this program.  If not, see <https://www.gnu.org/licenses/>.

''' This **private** module contains the commands that are called by the CLI.
'''
import numpy as np

from cartopy import crs as ccrs
from matplotlib import pyplot as plt

from pymagglobal import core
from pymagglobal import utils

from urllib.error import URLError
from warnings import warn


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


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 SimpleNamespace object returned by ArgumentParser.parse_args().

    Returns
    -------
    times : ndarray
        An equally spaced array in years.
    '''
    if args.begin is not None:
        if args.longterm:
            t_min = lt2yr(args.begin)
        else:
            t_min = args.begin
        args.model.valid_epoch(t_min)
    else:
        t_min = args.model.t_min

    if args.end is not None:
        if args.longterm:
            t_max = lt2yr(args.end)
        else:
            t_max = args.end
        args.model.valid_epoch(t_max)
    else:
        t_max = args.model.t_max

    if args.every is not None:
        times = np.arange(t_min, t_max + args.every, args.every)
    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]
    else:
        times = np.linspace(t_min, t_max, int(args.res))

    return times


def yr2lt(times):
    '''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).

    Examples
    --------
    >>> yr2lt(1950)
    0.0
    >>> yr2lt(-50)
    2.0
    '''
    return -(times - 1950) / 1000


def lt2yr(times):
    '''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
    '''
    return -times*1000 + 1950


def local_curve(args, fig=None):
    '''Handle the local command and create a local curve at a given
    location.

    Parameters
    ----------
    args : object
        The SimpleNamespace object returned by ArgumentParser.parse_args().
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.

    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
        local curve.
    '''
    # get an array of times
    times = args2times(args)
    # create a local curve using the core function, check is performed
    # in args2times
    curves = core.local_curve(times, (args.lat, args.lon), args.model,
                              field_type=args.type, check=False)
    # output formats for dif and nez components
    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

    if args.longterm:
        times = yr2lt(times)
    # 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'Local curves at ({args.lat}°, {args.lon}°) '
                          f'for {args.model.name}\n'
                          f'{args.t_label},'
                          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 local curve
    if not args.no_show or args.savefig is not None:
        if fig is None:
            fig = plt.figure(figsize=(10, 7))
        fig.suptitle(f'Local curves at ({args.lat}°, {args.lon}°) '
                     f'for {args.model.name}')
        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]))
            if args.longterm:
                axs[it].invert_xaxis()
                axs[it].set_xlabel(args.t_label)
            else:
                axs[it].set_xlabel(utils._nicelabel(args.t_label))
        fig.tight_layout(rect=[0, 0.03, 1, 0.95])
        return fig


def dipole_series(args, fig=None):
    '''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().
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.

    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.
    '''
    # get an array of times
    times = args2times(args)
    # 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)
    # if the --longterm flag is set, translate the years accordingly
    if args.longterm:
        times = yr2lt(times)

    # 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,
                   fmt=('%.2f', '%1.7e'),
                   delimiter=',',
                   header=f'Dipole moment series for '
                          f'{args.model.name}\n'
                          f'{args.t_label},m [m^2 A]')

    # if the --no-show flag is not set, plot the time series
    if not args.no_show or args.savefig is not None:
        if fig is None:
            fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111)
        ax.plot(times, dip_ser)
        ax.set_title(f'Dipole moment series for {args.model.name}')
        ax.set_ylabel(r'$m$ [$10^{22}$ m$^2$ A]')
        if args.longterm:
            ax.invert_xaxis()
            ax.set_xlabel(args.t_label)
        else:
            ax.set_xlabel(utils._nicelabel(args.t_label))
        fig.tight_layout()
        return fig


def coefficients_series(args):
    '''Handle the coeff-series command and provide the model time series for
    given coefficients.

    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
    _, _, coeffs = core.coefficients(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,
                                   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
    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'Coefficients series for {args.model.name}')
        ax.set_ylabel(r'$g_\ell^m$ [nT]')
        for ind, ell, m in zip(inds, args.degree, args.order):
            ax.plot(times, coeffs[:, 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


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


def coefficients_epoch(args):
    '''Handle the coeffs-epoch command and provide the model coefficients 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, coeffs = core.coefficients(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, coeffs]).T,
                   fmt=('%d', '%d', '%1.7e'),
                   delimiter=',',
                   header=f'Coefficients for {args.model.name} in nT at epoch '
                          f'{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(coeffs, ls, ms, args)


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$'):
    '''Plot Gauss coefficients.

    Parameters
    ----------
    gs : ndarray
        Gauss-coefficients.
    ls : ndarray
        SH degrees.
    ms : ndarray
        SH orders.
    unit : str, optional
        The unit of the coeffs. Default is nT.
    name : str, optional
        Headline name of what is being plotted.

    Returns
    -------
    matplotlib.figure.Figure
        matplotlib.figure.Figure object, containing a plot of the coefficients.
    '''
    ms = np.array(ms)
    lmax = np.max(ls)
    fig, axs = plt.subplots(ncols=1, nrows=lmax, sharex=True,
                            figsize=(10, 10),
                            subplot_kw={'xlim': (-lmax-0.1, lmax+0.1)})
    fig.suptitle(f'{name} for {args.model.name} at epoch '
                 f'{args.epoch} {args.t_unit}\n'
                 f'Units are {unit}.')
    for ell, ax in enumerate(axs, start=1):
        mask = np.equal(ls, ell)
        ax.scatter(ms[mask], gs[mask], marker='x')
        ymax = 1.1*np.abs(gs[mask]).max()
        ax.set_ylim(-ymax, ymax)
        ax.set_ylabel(fr'$\ell = {ell}$')
        ax.hlines(0, *ax.get_xlim(), color='gray', lw=0.5)
    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


def maps(args, fig=None):
    '''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().
    fig : matplotlib.figure.Figure, optional
        A figure to plot into. This is used for the GUI.

    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.
    '''
    # set up the points as expected by dsh_basis
    z_at = utils.get_z_at(args.res, t=args.epoch)
    # BUGFIX: white edges in cartopy
    z_at[1] -= 180
    # use the core function to get the field
    field = core.field(z_at, args.model, field_type=args.type)
    # 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 {args.model.name} at epoch '
                          f'{args.epoch} {args.t_unit}\n'
                          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 or args.savefig is not None:
        cmaps = ['RdBu', 'RdBu', 'RdBu']
        units = [r'$\mu$T', r'$\mu$T', r'$\mu$T']
        levels = [10, 10, 10]
        extends = ['neither', 'neither', 'neither']
        if args.type == 'dif':
            field[2] /= 1000

            levels[0] = [-40, -32, -24, -16, -8, 0, 8, 16, 24, 32, 40]
            levels[1] = np.linspace(-90, 90, 11)
            extends[0] = 'both'
            cmaps[2] = 'Blues'
            units[0] = r'deg.'
            units[1] = r'deg.'
        else:
            field /= 1000
            vmax = np.ceil(np.max(np.abs(field), axis=1))

            for it in range(3):
                levels[it] = np.linspace(-vmax[it], vmax[it], 11)

        cbar_hght = 0.08
        if fig is None:
            proj = ccrs.Mollweide()
            fig, axs = plt.subplots(1, 3, figsize=(13, 3.4),
                                    subplot_kw={'projection': proj})
            fig.tight_layout()
            fig.subplots_adjust(top=0.8, bottom=0.25, wspace=0.1)
            colaxs = []
            for it in range(3):
                bnds = axs[it].get_position().bounds

                colaxs.append(fig.add_axes([bnds[0],
                                            bnds[1]-0.06-cbar_hght,
                                            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
                colaxs.append(fig.add_axes([bnds[0],
                                            bnds[1]-0.02-0.5*cbar_hght,
                                            bnds[2],
                                            0.5*cbar_hght]))
            cbar_orientation = 'horizontal'

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

        fig.suptitle(f'Field maps for {args.model.name} at epoch '
                     f'{args.epoch} {args.t_unit}')

        for it in range(3):
            mappable = axs[it].tricontourf(plt_lat, plt_lon, field[it],
                                           cmap=cmaps[it],
                                           levels=levels[it],
                                           extend=extends[it])
            try:
                axs[it].coastlines(alpha=0.8, lw=0.5)
            except URLError:
                warn('Server connection timed out. Coastlines could not be '
                     'loaded.',
                     category=TimeOutWarning)
            axs[it].set_global()
            axs[it].set_title(f'{utils._names[args.type][it]} [{units[it]}]')
            colaxs[it].tick_params(labelsize=12)
            fig.colorbar(mappable,
                         cax=colaxs[it],
                         orientation=cbar_orientation)
        return fig