{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Polar Region Plots\n",
    "\n",
    "> Authors: Ashley Smith\n",
    ">\n",
    "> Abstract: Demonstrates access to (FAC, Te, Ne) measurements, and visualisation of them on Lon/Lat and MLT/QDLat plots."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See also:\n",
    "\n",
    " - https://github.com/pacesm/jupyter_notebooks/blob/master/AEBS/AEBS_AOB_FAC.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext watermark\n",
    "%watermark -i -v -p viresclient,pandas,xarray,matplotlib,cartopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from viresclient import SwarmRequest\n",
    "import datetime as dt\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Collection names:\")\n",
    "print(SwarmRequest().available_collections(\"FAC\", details=False))\n",
    "print(SwarmRequest().available_collections(\"IPD\", details=False), \"\\n\")\n",
    "print(\"FAC measurements:\\n\", SwarmRequest().available_measurements(\"FAC\"))\n",
    "print(\"IPD measurements:\\n\", SwarmRequest().available_measurements(\"IPD\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetch data separately from FAC and IPD collections over the same time window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set by IS0-8601 time:\n",
    "start = \"2018-05-01T00:00:00Z\"\n",
    "end = \"2018-05-01T12:00:00Z\"\n",
    "# Setting by datetime objects:\n",
    "# start = dt.datetime(2018, 5, 1, 0)\n",
    "# end = dt.datetime(2018, 5, 1, 12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_FACATMS_2F\")\n",
    "request.set_products(\n",
    "    measurements=[\"FAC\"],\n",
    "    auxiliaries=[\"MLT\", \"QDLat\", \"QDLon\"])\n",
    "data = request.get_between(start, end)\n",
    "ds_fac = data.as_xarray()\n",
    "ds_fac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_IPDAIRR_2F\")\n",
    "request.set_products(\n",
    "    measurements=[\"Ne\", \"Te\"],\n",
    "    auxiliaries=[\"MLT\", \"QDLat\", \"QDLon\"])\n",
    "data = request.get_between(start, end)\n",
    "ds_ipd = data.as_xarray()\n",
    "ds_ipd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why are some of the temperatures negative?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_ipd[\"Te\"].where(ds_ipd[\"Te\"] < 0, drop=True).head()  # NB only looking at first five entries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reorganise the data\n",
    "\n",
    "Note that the data are recorded at different sampling points:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_fac[\"Timestamp\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_ipd[\"Timestamp\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could have alternatively fetched the data together directly with `request.set_collection(\"SW_OPER_FACATMS_2F\", \"SW_OPER_IPDAIRR_2F\") ...` - in this case the server resolves the time discrepancy by using just the sampling times (and rate) of *the first collection given* (i.e. `\"SW_OPER_FACATMS_2F\"`) and interpolates the following collections onto the first time series with a nearest-neighbour method. This means that the `IPD` measurements would falsely be reported at the sampling times of the `FAC` measurements - this might not be a problem depending on your application, but we will avoid that issue here by accessing them as separate datasets.\n",
    "\n",
    "We can perform an outer join to merge the datasets into one object, where `nan`'s fill the empty sampling times in each input dataset. Let's also set the appropriate variables as coordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds_ipd.merge(ds_fac, join=\"outer\")\n",
    "coords = [\"Latitude\", \"Longitude\", \"Radius\", \"QDLat\", \"QDLon\", \"MLT\", \"Spacecraft\"]\n",
    "ds = ds.set_coords(coords)\n",
    "# NB: xarray merge does not handle the attributes\n",
    "#  so we must merge these manually\n",
    "ds = ds.assign_attrs({\"Sources\": ds_fac.Sources + ds_ipd.Sources})\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use this single object, `ds`, to access data in the rest of the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualising the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualise the measurements using the shortcut xarray plotting methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.plot.scatter(x=\"Longitude\", y=\"Latitude\", hue=\"FAC\", s=0.2, robust=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.plot.scatter(x=\"Longitude\", y=\"Latitude\", hue=\"Ne\", s=0.2, robust=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.plot.scatter(x=\"Longitude\", y=\"Latitude\", hue=\"Te\", s=0.2, robust=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.plot.scatter(x=\"MLT\", y=\"QDLat\", hue=\"Te\", s=0.2, robust=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geographic (Lon/Lat) and MLT/QDLat plots\n",
    "\n",
    "To make more complex figures, it is usually necessary to drop down to the lower level matplotlib interface.\n",
    "\n",
    "First let's set up figures onto which data can be plotted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    \n",
    "def create_axes_geo(title=\"\", figsize=(10, 5)):\n",
    "    \"\"\"Create an empty geographic figure with North/South views\n",
    "    \n",
    "    Args:\n",
    "        title (str)\n",
    "        figsize (tuple): (width, height)\n",
    "        \n",
    "    Returns:\n",
    "        matplotlib.figure.Figure, [GeoAxesSubplot, GeoAxesSubplot]\n",
    "    \"\"\"\n",
    "    fig = plt.figure(figsize=figsize)\n",
    "    fig.suptitle(title, fontsize=15)\n",
    "    # Geographic (lat/lon) views\n",
    "    ax_N_geo = plt.subplot2grid(\n",
    "        (1, 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_geo = plt.subplot2grid(\n",
    "        (1, 2), (0, 1),\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_N_geo.set_extent([-180, 180, 50, 90], ccrs.PlateCarree())\n",
    "    ax_S_geo.set_extent([-180, 180, -90, -50], ccrs.PlateCarree())\n",
    "    for ax in [ax_N_geo, ax_S_geo]:\n",
    "        _apply_circular_boundary(ax)\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_N_geo.gridlines(ylocs=[50, 60, 70, 80, 90])\n",
    "    ax_S_geo.gridlines(ylocs=[-90, -80, -70, -60, -50])\n",
    "    ax_N_geo.set_title(\"North\")\n",
    "    ax_S_geo.set_title(\"South\")\n",
    "    return fig, [ax_N_geo, ax_S_geo]\n",
    "\n",
    "def create_axes_mlt(title=\"\", figsize=(10, 5)):\n",
    "    \"\"\"Create an empty MLT/QDLat figure with North/South views\n",
    "    \n",
    "    Args:\n",
    "        title (str)\n",
    "        figsize (tuple): (width, height)\n",
    "        \n",
    "    Returns:\n",
    "        matplotlib.figure.Figure, [PolarAxesSubplot, PolarAxesSubplot]\n",
    "    \"\"\"\n",
    "    fig = plt.figure(figsize=figsize)\n",
    "    fig.suptitle(title, fontsize=15)\n",
    "    # QDLat/MLT views\n",
    "    ax_N_mlt = plt.subplot2grid(\n",
    "        (1, 2), (0, 0),\n",
    "        projection=\"polar\"\n",
    "    )\n",
    "    ax_S_mlt = plt.subplot2grid(\n",
    "        (1, 2), (0, 1),\n",
    "        projection=\"polar\"\n",
    "    )\n",
    "    for ax in [ax_N_mlt, ax_S_mlt]:\n",
    "        ax.set_ylim(0, 40)\n",
    "        ax.set_yticks([0, 10, 20, 30, 40])\n",
    "        ax.set_xticklabels([\"%2.2i\" % (x*12/np.pi) for x in ax.get_xticks()])\n",
    "        ax.set_theta_zero_location(\"S\")\n",
    "        ax.set_yticklabels([])\n",
    "    ax_N_mlt.set_title(\"North\")\n",
    "    ax_S_mlt.set_title(\"South\")\n",
    "    return fig, [ax_N_mlt, ax_S_mlt]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These can be used together with the xarray plotting methods which will make some decisions for us automatically (like setting up the colorbars):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = create_axes_geo()\n",
    "for ax in axes:\n",
    "    ds.plot.scatter(x=\"Longitude\", y=\"Latitude\", hue=\"Te\", s=0.1,\n",
    "                    ax=ax, transform=ccrs.PlateCarree(), robust=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(the MLT plot case is more complicated because of the transformations needed to plot data onto the polar projections)\n",
    "\n",
    "To give more control over the plotting, we use matplotlib directly:\n",
    "\n",
    "### Functions to help plot onto the axes specified above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _subselect(ds, var, subselect_factor):\n",
    "    _ds = ds.where(~np.isnan(ds[var]), drop=True)\n",
    "    _ds = _ds.isel({\"Timestamp\": slice(0, -1, subselect_factor)})\n",
    "    return _ds\n",
    "\n",
    "def plot_geo(ax, ds, hemisphere=\"north\", **kwargs):\n",
    "    \"\"\"\n",
    "    kwargs must contain \"var\" to plot from ds\n",
    "    \n",
    "    Args:\n",
    "        ax (matplotlib.axes)\n",
    "        ds (xarray.Dataset)\n",
    "        hemisphere (str): \"north\" or \"south\"\n",
    "    \"\"\"\n",
    "    # Identify data variable to plot\n",
    "    var = kwargs.pop(\"var\")\n",
    "    # Sub-select data by a given factor\n",
    "    subselect = kwargs.pop(\"subselect\", None)\n",
    "    if subselect:\n",
    "        _ds = _subselect(ds, var, subselect)\n",
    "    else:\n",
    "        _ds = ds\n",
    "    # Select only either NH or SH data\n",
    "    if hemisphere == \"north\":\n",
    "        _ds = _ds.where(ds[\"Latitude\"] > 50, drop=True)\n",
    "    elif hemisphere == \"south\":\n",
    "        _ds = _ds.where(ds[\"Latitude\"] < -50, drop=True)\n",
    "    # Do the actual plotting\n",
    "    p = ax.scatter(_ds[\"Longitude\"], _ds[\"Latitude\"], c=_ds[var],\n",
    "                   transform=ccrs.PlateCarree(), **kwargs)\n",
    "    return p\n",
    "\n",
    "def plot_mlt(ax, ds, hemisphere=\"north\", **kwargs):\n",
    "    \"\"\"\n",
    "    kwargs must contain \"var\" to plot from ds\n",
    "    \n",
    "    Args:\n",
    "        ax (matplotlib.axes)\n",
    "        ds (xarray.Dataset)\n",
    "        hemisphere (str): \"north\" or \"south\"\n",
    "    \"\"\"\n",
    "    # Identify data variable to plot\n",
    "    var = kwargs.pop(\"var\")\n",
    "    # Sub-select data by a given factor\n",
    "    subselect = kwargs.pop(\"subselect\", None)\n",
    "    if subselect:\n",
    "        _ds = _subselect(ds, var, subselect)\n",
    "    else:\n",
    "        _ds = ds\n",
    "    # Select only either NH or SH data\n",
    "    if hemisphere == \"north\":\n",
    "        _ds = _ds.where(ds[\"QDLat\"] > 50, drop=True)\n",
    "        h = 1\n",
    "    elif hemisphere == \"south\":\n",
    "        _ds = _ds.where(ds[\"QDLat\"] < -50, drop=True)\n",
    "        h = -1\n",
    "    # Transformation to the polar coordinates\n",
    "    def _scatter(x, y, *args, **kwargs):\n",
    "        return ax.scatter(x*(np.pi/12), 90 - y*h, *args, **kwargs)\n",
    "    p = _scatter(_ds[\"MLT\"], _ds[\"QDLat\"], c=_ds[var],\n",
    "                 **kwargs)\n",
    "    return p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Te on both GEO and MLT views"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the bare figures\n",
    "fig_geo, (ax_N_geo, ax_S_geo) = create_axes_geo()\n",
    "fig_mlt, (ax_N_mlt, ax_S_mlt) = create_axes_mlt()\n",
    "\n",
    "# Set the parameters to control the plotting\n",
    "options = dict(\n",
    "    var=\"Te\",\n",
    "    cmap=mpl.cm.plasma,\n",
    "    norm=mpl.colors.Normalize(vmin=3e3, vmax=6e3),\n",
    "    s=0.1,\n",
    ")\n",
    "\n",
    "# Plot onto each axes\n",
    "plot_geo(ax_N_geo, ds, hemisphere=\"north\", **options)\n",
    "plot_geo(ax_S_geo, ds, hemisphere=\"south\", **options)\n",
    "plot_mlt(ax_N_mlt, ds, hemisphere=\"north\", **options)\n",
    "plot_mlt(ax_S_mlt, ds, hemisphere=\"south\", **options)\n",
    "\n",
    "# Add colorbars\n",
    "for fig in [fig_geo, fig_mlt]:\n",
    "    cax = fig.add_axes([0.4, 0.15, 0.2, 0.02])\n",
    "    var = options[\"var\"]\n",
    "    label = f\"{var} [{ds[var].units}]\"\n",
    "    cbar = mpl.colorbar.ColorbarBase(cax,\n",
    "                                     cmap=options[\"cmap\"],\n",
    "                                     norm=options[\"norm\"],\n",
    "                                     orientation='horizontal',\n",
    "    #                                  ticks=[norm.vmin, (norm.vmax+norm.vmin)/2, norm.vmax],\n",
    "                                     label=label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot both FAC and Ne on one figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up title based on approx start/end times\n",
    "time1 = ds[\"Timestamp\"][0].values\n",
    "time2 = ds[\"Timestamp\"][-1].values\n",
    "def format_dt64(dt64, form=\"%Y-%m-%d %H:%M\"):\n",
    "    # Convert numpy.datetime64 [ns] -> datetime\n",
    "    time = dt.datetime.utcfromtimestamp(dt64.astype(int) * 1e-9)\n",
    "    return time.strftime(form)\n",
    "spacecraft = ds[\"Spacecraft\"][0].values\n",
    "title = f\"{format_dt64(time1)}\\n{format_dt64(time2)}\\n(Swarm {spacecraft})\"\n",
    "\n",
    "fig, (ax_N_mlt, ax_S_mlt) = create_axes_mlt(title=title)\n",
    "\n",
    "options_Ne = dict(\n",
    "    var=\"Ne\",\n",
    "    cmap=mpl.cm.viridis,\n",
    "    norm=mpl.colors.Normalize(vmin=10e3, vmax=100e3),\n",
    "    s=0.1, zorder=2,\n",
    ")\n",
    "# - scatterplot point size (s) and data rate (subselect)\n",
    "#     need to be balanced to make them legible\n",
    "# - zorder controls the layers of plots (which one goes on top)\n",
    "# - Try a horizontal line as the marker\n",
    "#     (only works when the satellite tracks run vertically)\n",
    "#     ref: https://matplotlib.org/3.1.0/api/markers_api.html\n",
    "options_FAC = dict(\n",
    "    var=\"FAC\", subselect=10,\n",
    "    cmap=mpl.cm.seismic,\n",
    "    norm=mpl.colors.SymLogNorm(linthresh=1, linscale=1, vmin=-20,vmax=20),\n",
    "    s=80, zorder=1, marker=\"_\",\n",
    ")\n",
    "\n",
    "for options in [options_FAC, options_Ne]:\n",
    "    plot_mlt(ax_N_mlt, ds, hemisphere=\"north\", **options)\n",
    "    plot_mlt(ax_S_mlt, ds, hemisphere=\"south\", **options)\n",
    "\n",
    "# Set colorbar locations\n",
    "cax1 = fig.add_axes([0.4, 0.15, 0.2, 0.02])\n",
    "cax2 = fig.add_axes([0.4, 0.0, 0.2, 0.02])\n",
    "# Draw colorbars\n",
    "for cax, options in zip([cax1, cax2], [options_FAC, options_Ne]):\n",
    "    var = options[\"var\"]\n",
    "    label = f\"{var} [{ds[var].units}]\"\n",
    "    cbar = mpl.colorbar.ColorbarBase(cax,\n",
    "                                     cmap=options[\"cmap\"],\n",
    "                                     norm=options[\"norm\"],\n",
    "                                     orientation='horizontal',\n",
    "                                     label=label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "next steps:\n",
    "\n",
    "- visualisations of auroral oval boundaries etc, from AEBS products:\n",
    " - https://github.com/pacesm/jupyter_notebooks/blob/master/AEBS/AEBS_AOB_FAC.ipynb\n",
    "- use `cartopy.feature.nightshade` to show terminator (for shorter time periods)\n",
    " - https://scitools.org.uk/cartopy/docs/latest/gallery/aurora_forecast.html\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}