{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the python backend"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we illustrate some ways of using `pymagglobal` from within python. For explicit documentation, see [here](https://sec23.git-pages.gfz-potsdam.de/korte/pymagglobal/package_documentation.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a `Model` instance\n",
    "\n",
    "The first step to work with `pymagglobal` models is to make them accessible by instancing the `Model` class. We therefore import `Model` and create an instance by passing the name of the model we want to access, in this case `GGFMB`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymagglobal import Model\n",
    "\n",
    "ggfmb = Model('GGFMB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can access some information about the model, for example the maximum degree:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ggfmb.l_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the `valid_epoch` method to check if an epoch is covered by the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass the epoch as yr.\n",
    "\n",
    "ggfmb.valid_epoch(-700_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " `valid_epoch` returns the input, but throws a warning if the input is not covered by the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass the epoch as yr.\n",
    "\n",
    "ggfmb.valid_epoch(-2_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using `core` functions\n",
    "\n",
    "The most straight-forward way of extracting specific information from a model is to use the `core` functions of `pymagglobal`. Check the [documentation](https://sec23.git-pages.gfz-potsdam.de/korte/pymagglobal/package_documentation.html) to see which functions are available.\n",
    "\n",
    "Let's assume we want to get a local time series of the inclination at a specific location. We can use the `local_curve` function for that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from pymagglobal import local_curve\n",
    "\n",
    "# Honululu as lat, lon tuple\n",
    "loc = (21.306944, -157.858333)\n",
    "\n",
    "# Create an evenly spaced array over the ggfmb interval\n",
    "\n",
    "times = np.linspace(\n",
    "    ggfmb.t_min,\n",
    "    ggfmb.t_max,\n",
    "    501,\n",
    ")\n",
    "d, i, f = local_curve(times, loc, ggfmb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now plot the result. Note that we convert the input to kiloyears before present for better readability:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from pymagglobal.utils import yr2lt\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
    "\n",
    "ax.plot(\n",
    "    yr2lt(times),\n",
    "    i,\n",
    ")\n",
    "ax.invert_xaxis()\n",
    "ax.set_xlabel('time [ka BP]')\n",
    "ax.set_ylabel('inclination [deg.]');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We might also be interested in a map of the radial component (Z) of the field at the Earth's surface. We can use the `field` method to create such a map:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymagglobal import field\n",
    "from pymagglobal.utils import get_grid, lt2yr\n",
    "\n",
    "# Set up a grid at the Earth's surface\n",
    "# The epoch 780 ka BP is transformed to years using lt2yr\n",
    "grid = get_grid(1000, t=lt2yr(780))\n",
    "# for plotting purposes shift the array\n",
    "grid[1] -= 180\n",
    "# access is then straight-forward\n",
    "n, e, z = field(grid, ggfmb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the result using `cartopy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cartopy import crs as ccrs\n",
    "\n",
    "proj = ccrs.Mollweide()\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5), subplot_kw={'projection': proj})\n",
    "\n",
    "plt_lat, plt_lon, _ = proj.transform_points(\n",
    "    ccrs.Geodetic(),\n",
    "    grid[1],\n",
    "    90-grid[0],\n",
    ").T\n",
    "\n",
    "vmax = np.max(np.abs(z))\n",
    "ax.tricontourf(\n",
    "    plt_lat,\n",
    "    plt_lon,\n",
    "    z,\n",
    "    vmax=vmax,\n",
    "    vmin=-vmax,\n",
    "    cmap='RdBu',\n",
    ")\n",
    "ax.set_global()\n",
    "ax.coastlines();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using model coefficients directly\n",
    "\n",
    "If we are intersted in something that is not available directly via the `core` functions, we might be able to use model coefficients directly and calculate the quantity of interst. Here we show how to generate a map of the radial (Z) component of the field at the core-mantle boundary. There are two ways to do this. Either we scale the coefficients and use unscaled basis functions, or we might also have the basis functions include a scaling. To highlight several aspects of `pymagglobal` we pursue the first approach. Initially, we access the coefficients for the specific epoch at the Earth's surface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymagglobal import coefficients\n",
    "\n",
    "l, m, coeffs = coefficients(lt2yr(780), ggfmb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we scale the coeffients to the core-mantle boundary (3480 km):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymagglobal.utils import scaling, REARTH\n",
    "\n",
    "coeffs_cmb = coeffs * scaling(REARTH, 3480, ggfmb.l_max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to calculate the field components, we evaluate the basis functions at the grid specified before. The basis functions are derivatives of the spherical harmonics. The function `dsh_basis` takes array of shape (L, 3*N), where L is the number of Gauss coefficients, as third argument and fills it with the appropriate values. We calculate the number of coeffients from the maximum spherical harmonics degree using `lmax2N`. Note that by default, `dsh_basis` assumes that the coefficients are given at the Earth's surface. I.e. if the grid is also at the Earth's surface, no scaling is applied. If however the radial grid entries `grid[2]` differ from `REARTH`, `dsh_basis` automatically includes a scaling. Check the end of this notebook for the way to handle this scalings explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymagglobal.utils import dsh_basis, lmax2N\n",
    "\n",
    "base = np.empty((lmax2N(ggfmb.l_max), 3*grid.shape[1]))\n",
    "dsh_basis(ggfmb.l_max, grid, base)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculating the field is then straight forward, we simply multiply the coefficients and the basis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nez_cmb = coeffs_cmb @ base"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To access the radial component we reshape the field array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nez_cmb = nez_cmb.reshape(grid.shape[1], 3).T\n",
    "z_cmb = nez_cmb[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can again plot the results using cartopy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5), subplot_kw={'projection': proj})\n",
    "\n",
    "vmax = np.max(np.abs(z_cmb))\n",
    "ax.tricontourf(\n",
    "    plt_lat,\n",
    "    plt_lon,\n",
    "    z_cmb,\n",
    "    vmax=vmax,\n",
    "    vmin=-vmax,\n",
    "    cmap='RdBu',\n",
    ")\n",
    "ax.set_global()\n",
    "ax.coastlines();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case we could have also set the radius of the grid to the core-mantle boundary and use the `field` routine. This implicitly moves the scaling to the basis functions as mentioned above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_cmb = np.copy(grid)\n",
    "grid_cmb[2] = 3480\n",
    "\n",
    "_, _, z_cmb_2 = field(grid_cmb, ggfmb)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5), subplot_kw={'projection': proj})\n",
    "\n",
    "ax.tricontourf(\n",
    "    plt_lat,\n",
    "    plt_lon,\n",
    "    z_cmb_2,\n",
    "    vmax=vmax,\n",
    "    vmin=-vmax,\n",
    "    cmap='RdBu',\n",
    ")\n",
    "ax.set_global()\n",
    "ax.coastlines();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The formally correct/explicit way of using `dsh_basis` like above would be to both move the grid to the core-mantle boundary and use `dsh_basis` with 3480 km as reference radius like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_cmb = np.empty((lmax2N(ggfmb.l_max), 3*grid_cmb.shape[1]))\n",
    "dsh_basis(ggfmb.l_max, grid_cmb, base_cmb, R=3480)\n",
    "\n",
    "nez_cmb_3 = coeffs_cmb @ base_cmb\n",
    "nez_cmb_3 = nez_cmb_3.reshape(grid.shape[1], 3).T\n",
    "z_cmb_3 = nez_cmb_3[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5), subplot_kw={'projection': proj})\n",
    "\n",
    "ax.tricontourf(\n",
    "    plt_lat,\n",
    "    plt_lon,\n",
    "    z_cmb_3,\n",
    "    vmax=vmax,\n",
    "    vmin=-vmax,\n",
    "    cmap='RdBu',\n",
    ")\n",
    "ax.set_global()\n",
    "ax.coastlines();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the implicit scaling applied above, `base_cmb` and `base` are actually identical:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert(np.all(base == base_cmb))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pymagglobal",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}