{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geomagnetic Models (eoxmagmod + contours)\n",
    "\n",
    "> Abstract: Demonstrate basic usage of eoxmagmod for forwards evaluation of magnetic model, together with contour visualisations with cartopy.\n",
    "\n",
    "[eoxmagmod](https://github.com/ESA-VirES/MagneticModel/) is used internally within VirES to perform the forward evaluations of geomagnetic models. This notebook is a demonstration of using it directly for the simple .shc file case, though it can also be used to evaluate the more complex models. There is not much documentation but we could work to improve this (and the installation and usability) if it is useful.\n",
    "\n",
    "Here we show a basic setup of contour plots using cartopy. It would be good to build a richer set of preset visualisations as part of an importable library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See also:\n",
    "\n",
    " - https://github.com/pacesm/jupyter_notebooks/blob/master/examples/CHAOS-6_Cartopy_Contours.ipynb\n",
    " - https://github.com/MagneticEarth/MagneticEarth_notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "if 'CI' in os.environ:\n",
    "    os.environ['CDF_LIB'] = '/srv/conda/envs/notebook/lib'\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark -i -v -p viresclient,pandas,xarray,matplotlib,scipy,cartopy,eoxmagmod,chaosmagpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import datetime as dt\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.interpolate import interp1d\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import eoxmagmod\n",
    "from chaosmagpy.plot_utils import nio_colormap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions to do the legwork of evaluation and plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grid(nlats=180, nlons=360):\n",
    "    \"\"\"A global grid of the specified size\"\"\"\n",
    "    _lats = np.linspace(-90, 90, nlats + 1)\n",
    "    _lons = np.linspace(-180, 180, nlons + 1)\n",
    "    coords = np.empty((_lats.size, _lons.size, 3))\n",
    "    coords[:,:,1], coords[:,:,0] = np.meshgrid(_lons, _lats)\n",
    "    coords[:,:,2] = 0 # height above WGS84 in km\n",
    "    return coords\n",
    "\n",
    "def eval_model(time=dt.datetime(2020, 1, 1), coords=grid(),\n",
    "               shc_model=eoxmagmod.data.CHAOS_CORE_LATEST):\n",
    "    \"\"\"Evaluate a .shc model at a fixed time\n",
    "\n",
    "    Args:\n",
    "        time (datetime)\n",
    "        coords (ndarray)\n",
    "        shc_model (str): path to file\n",
    "\n",
    "    Returns:\n",
    "        dict: magnetic field vector components:\n",
    "            https://intermagnet.github.io/faq/10.geomagnetic-comp.html\n",
    "    \"\"\"\n",
    "    # Convert Python datetime to MJD2000\n",
    "    epoch = eoxmagmod.time_util.datetime_to_decimal_year(time)\n",
    "    mjd2000 = eoxmagmod.decimal_year_to_mjd2000(epoch)\n",
    "    # Load model\n",
    "    model = eoxmagmod.load_model_shc(shc_model)\n",
    "    # Evaluate in North, East, Up coordinates\n",
    "    height = eoxmagmod.GEODETIC_ABOVE_WGS84\n",
    "    b_neu = model.eval(mjd2000, coords, height, height)\n",
    "    # Inclination (I), declination (D), intensity (F)\n",
    "    inc, dec, F = eoxmagmod.vincdecnorm(b_neu)\n",
    "    return {\"X\": b_neu[:,:,0], \"Y\": b_neu[:,:,1], \"Z\": -b_neu[:,:,2],\n",
    "            \"I\": inc, \"D\":dec, \"F\":F}\n",
    "\n",
    "def _plot_contours(ax, x, y, z, *args, **kwargs):\n",
    "    transform_before_plotting = kwargs.pop(\"transform_before_plotting\", False)\n",
    "    if transform_before_plotting:\n",
    "        # transform coordinates *before* passing them to the plotting function\n",
    "        tmp = ax.projection.transform_points(ccrs.PlateCarree(), x, y)\n",
    "        x_t, y_t = tmp[..., 0], tmp[..., 1]\n",
    "        return ax.contour(x_t, y_t, z, *args, **kwargs)\n",
    "    else:\n",
    "        # ... transformation performed by the plotting function creates glitches at the antemeridian\n",
    "        kwargs[\"transform\"] = ccrs.PlateCarree()\n",
    "        return ax.contour(x, y, z, *args, **kwargs)\n",
    "\n",
    "def plot_contours(ax, x, y, z, *args, **kwargs):\n",
    "    fmt = kwargs.pop(\"fmt\", \"%g\")\n",
    "    fontsize = kwargs.pop(\"fontsize\", 6)\n",
    "    ax.add_feature(cfeature.LAND, facecolor=(1.0, 1.0, 0.9))\n",
    "    ax.add_feature(cfeature.OCEAN, facecolor=(0.9, 1.0, 1.0))\n",
    "    ax.add_feature(cfeature.COASTLINE, edgecolor='silver')\n",
    "    ax.gridlines()\n",
    "    cs = _plot_contours(ax, x, y, z, *args, **kwargs)\n",
    "    ax.clabel(cs, cs.levels, inline=True, fmt=fmt, fontsize=fontsize)\n",
    "\n",
    "def _apply_circular_boundary(ax):\n",
    "    \"\"\"Make cartopy axes have round borders.\n",
    "    See https://scitools.org.uk/cartopy/docs/v0.15/examples/always_circular_stereo.html\n",
    "    \n",
    "    Notes:\n",
    "        Inline contour labels are still appearing outside the boundary\n",
    "    \"\"\"\n",
    "    theta = np.linspace(0, 2*np.pi, 100)\n",
    "    center, radius = [0.5, 0.5], 0.5\n",
    "    verts = np.vstack([np.sin(theta), np.cos(theta)]).T\n",
    "    circle = mpl.path.Path(verts * radius + center)\n",
    "    ax.set_boundary(circle, transform=ax.transAxes)\n",
    "\n",
    "def contours_NorthSouthMoll(coords, data, units, title, **options):\n",
    "    \"\"\"Generate trio of AzimuthalEquidistant (N/S) and Mollweide projections\n",
    "    \n",
    "    Args:\n",
    "        coords (ndarray)\n",
    "        data (ndarray)\n",
    "        units (str)\n",
    "        title (str)\n",
    "    \"\"\"\n",
    "    # Set up figure with North/South/Mollweide maps\n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    fig.suptitle(title, fontsize=15)\n",
    "    ax_N = plt.subplot2grid(\n",
    "        (2, 2), (0, 0),\n",
    "        projection=ccrs.AzimuthalEquidistant(\n",
    "            central_longitude=0.0, central_latitude=90.0,\n",
    "            false_easting=0.0, false_northing=0.0, globe=None\n",
    "        )\n",
    "    )\n",
    "    ax_S = plt.subplot2grid(\n",
    "        (2, 2), (0, 1), colspan=2,\n",
    "        projection=ccrs.AzimuthalEquidistant(\n",
    "            central_longitude=0.0, central_latitude=-90.0,\n",
    "            false_easting=0.0, false_northing=0.0, globe=None\n",
    "        )\n",
    "    )\n",
    "    ax_Moll = plt.subplot2grid(\n",
    "        (2, 2), (1, 0), colspan=2,\n",
    "        projection=ccrs.Mollweide()\n",
    "    )\n",
    "    ax_N.set_extent([-180, 180, 40, 90], ccrs.PlateCarree())\n",
    "    ax_S.set_extent([-180, 180, -90, -40], ccrs.PlateCarree())\n",
    "#     _apply_circular_boundary(ax_N)\n",
    "#     _apply_circular_boundary(ax_S)\n",
    "    # Plot on the contours from the given data\n",
    "    # Set options and update with any overrides provided\n",
    "    kwargs = {\n",
    "        \"fmt\": f\"%g {units}\",\n",
    "        \"linewidths\": 2,\n",
    "        \"fontsize\": 10\n",
    "    }\n",
    "    kwargs.update(options)\n",
    "    ## Masks to select the northern and southern parts separately\n",
    "    # Notes:\n",
    "    #  - the polar orthographic projection works only with all points\n",
    "    #    within the maps extent (i.e., visible part of the globe).\n",
    "    north = (coords[:, 0, 0] > 0).nonzero()[0]\n",
    "    south = (coords[:, 0, 0] < 0).nonzero()[0]\n",
    "    # Northern and Southern Hemipshere plots\n",
    "    plot_contours(ax_N,\n",
    "                  coords[north, :, 1], coords[north, :, 0], data[north, :],\n",
    "                  transform_before_plotting=True, **kwargs)\n",
    "    plot_contours(ax_S,\n",
    "                  coords[south, :, 1], coords[south, :, 0], data[south, :],\n",
    "                  transform_before_plotting=True, **kwargs)\n",
    "    # Mollweide plot\n",
    "    plot_contours(ax_Moll,\n",
    "                  coords[..., 1], coords[..., 0], data,\n",
    "                  transform_before_plotting=False, **kwargs)\n",
    "#     fig.tight_layout()\n",
    "    fig.subplots_adjust(hspace=0.05, wspace=0.05)\n",
    "    return fig, [ax_N, ax_S, ax_Moll]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform the model evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = dt.datetime(2020, 1, 1)\n",
    "# You can supply your own shc file here\n",
    "#  Just set shc_model as the file path\n",
    "shc_model = eoxmagmod.data.CHAOS_CORE_LATEST\n",
    "coords = grid()\n",
    "mag_components = eval_model(t, coords, shc_model)\n",
    "mag_components.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The evaluated model vector components are stored as a dictionary, `mag_components`, so that you can access the `X` (Northwards) component as `mag_components['X']` and so on, together with the matching coordinates stored in `coords`.\n",
    "\n",
    "X: Northward, Y: Eastward, Z: Downward, I: Inclination, D: Declination, F: intensity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contours_NorthSouthMoll(\n",
    "    coords, mag_components[\"F\"], \"nT\",\n",
    "    f\"Intensity\\n{t.date()}\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Options can be fed through to the underlying matplotlib calls\n",
    "options = {\n",
    "    \"levels\": [-180, -135, -90, -60, -30, -20, -15, -10, -5, 0, 5, 10, 15, 20, 30, 60, 90, 135, 180],\n",
    "    \"cmap\": \"twilight\"\n",
    "}\n",
    "contours_NorthSouthMoll(\n",
    "    coords, mag_components[\"D\"], \"°\",\n",
    "    f\"Declination\\n{t.date()}\", **options);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contours_NorthSouthMoll(\n",
    "    coords, mag_components[\"Z\"], \"nT\",\n",
    "    f\"Downward component\\n{t.date()}\",\n",
    "    levels=20, cmap=nio_colormap());"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}