{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geomagnetic Model Residuals\n",
    "\n",
    "> Abstract: An exploration of geomagnetic data-model residuals evaluated through VirES (along Swarm orbits)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext watermark\n",
    "%watermark -i -v -p viresclient,pandas,xarray,matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from viresclient import SwarmRequest\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One day of residuals\n",
    "\n",
    "(data minus a core field model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_MAGA_LR_1B\")\n",
    "request.set_products(\n",
    "    measurements=[\"B_NEC\"],\n",
    "    models=[\"CHAOS-Core\"],\n",
    "    residuals=True,\n",
    "    sampling_step=\"PT5S\",\n",
    "    auxiliaries=[\"MLT\", \"QDLat\", \"SunZenithAngle\"]\n",
    ")\n",
    "data = request.get_between(\n",
    "    \"2019-01-01\", \"2019-01-02\",\n",
    "    asynchronous=False, show_progress=False\n",
    ")\n",
    "ds = data.as_xarray()\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "xarray provides some convenient direct plotting tools through integration with matplotlib:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"B_NEC_res_CHAOS-Core\"].plot.line(x=\"Timestamp\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This allows some very compact plotting commands to create complex figures, if you know how! This kind of plotting takes a while to learn and experiment with to get right, but results in short and manageable code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "facetgrid = ds.plot.scatter(x=\"QDLat\", y=\"B_NEC_res_CHAOS-Core\", col=\"NEC\",\n",
    "                s=1, hue=\"SunZenithAngle\", cmap=\"hot_r\", alpha=0.6, linewidths=0)\n",
    "for ax in facetgrid.axs.flat:\n",
    "    ax.set_facecolor(\"grey\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(showing that we are in daylight more in the Southern hemisphere, while in darkness in the Northern, because the data is from January - Northern winter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "facetgrid = ds.plot.scatter(x=\"QDLat\", y=\"B_NEC_res_CHAOS-Core\", col=\"NEC\",\n",
    "                s=1, hue=\"MLT\", cmap=\"twilight_shifted\", alpha=0.6, linewidths=0)\n",
    "for ax in facetgrid.axs.flat:\n",
    "    ax.set_facecolor(\"grey\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(showing the rapid movement through local time sectors (MLT) when the satellite passes through polar regions)\n",
    "\n",
    "The above figures show large residuals in the polar regions due to currents in the auroral oval. These cause greater disturbance in the Northward (N) and Eastward (E) components of the magnetic field because of the geometry of the field-aligned currents (FACs) which cause the magnetic disturbance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Access 1 year of data and models\n",
    "\n",
    "Let's inspect a whole year of data (at 2-minute sampling):\n",
    "\n",
    "- The `B_NEC` vector from Swarm Alpha\n",
    "- All parts of the CHAOS model (core, crust, magnetosphere) evaluated at the same points\n",
    "- MLT, QDLat, and SunZenithAngle which are also evaluated on-the-fly\n",
    "- Filtered to remove some bad data (according to `Flags_F`) and restricted to geomagnetically quiet data (according to `Kp`)\n",
    "\n",
    "We fetch the model values themselves as well as the measurements (instead of just the residuals as above) so that we can manipulate all the different components locally.\n",
    "\n",
    "**Warning**: This will take several minutes to process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_MAGA_LR_1B\")\n",
    "request.set_products(\n",
    "    measurements=[\"B_NEC\"],\n",
    "    models=[\"CHAOS-Core\", \"CHAOS-Static\", \"CHAOS-MMA-Primary\", \"CHAOS-MMA-Secondary\"],\n",
    "    auxiliaries=[\"MLT\", \"QDLat\", \"SunZenithAngle\", \"OrbitNumber\"],\n",
    "    sampling_step=\"PT120S\"\n",
    ")\n",
    "request.set_range_filter(\"Flags_F\", 0, 1)\n",
    "request.set_range_filter(\"Kp\", 0, 3)\n",
    "# request.set_range_filter(\"SunZenithAngle\", 100, 180)\n",
    "data = request.get_between(\"2019-01-01\", \"2020-01-01\")\n",
    "ds = data.as_xarray()\n",
    "# Remove the extra-long Sources list\n",
    "ds.attrs[\"Sources\"] = \"Many\"\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can construct `xarray.DataArray`'s based on `ds`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"B_NEC\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"B_NEC\"] - ds[\"B_NEC_CHAOS-Core\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and we can plot these directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(ds[\"B_NEC\"] - ds[\"B_NEC_CHAOS-Core\"]).plot.line(x=\"Timestamp\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can slice out a particular time window:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Selects two days\n",
    "ds.sel({\"Timestamp\": slice(\"2019-01-01\", \"2019-01-02\")})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and subselect according to parts of the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.where(ds[\"SunZenithAngle\"] > 100, drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the above to construct a plot based on part of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select one day\n",
    "_ds = ds.sel({\"Timestamp\": slice(\"2019-01-01\", \"2019-01-02\")})\n",
    "# Select nightside data from it\n",
    "_ds_dark = _ds.where(_ds[\"SunZenithAngle\"] > 100, drop=True)\n",
    "# Append a custom residual of B-MCO-MLI-MMA\n",
    "_ds_dark[\"B_res_full\"] = _ds_dark[\"B_NEC\"] - _ds_dark[\"B_NEC_CHAOS-Core\"] \\\n",
    "                                           - _ds_dark[\"B_NEC_CHAOS-Static\"] \\\n",
    "                                           - _ds_dark[\"B_NEC_CHAOS-MMA-Primary\"] \\\n",
    "                                           - _ds_dark[\"B_NEC_CHAOS-MMA-Secondary\"]\n",
    "_ds_dark.plot.scatter(x=\"QDLat\", y=\"B_res_full\", hue=\"OrbitNumber\", col=\"NEC\",\n",
    "                      cmap=\"viridis\", s=1, linewidths=0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Residuals to model combinations (one month)\n",
    "\n",
    "First we adjust the `Dataset` so that it contains the custom residuals themselves. We store them as both regular `data variables` and as a higher dimensional `\"B_residuals\"` which includes all of them. This is inefficient but is convenient for the plotting tools used below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assign_residuals(ds):\n",
    "    # Work on a copy of the Dataset so we don't disturb the original\n",
    "    ds = ds.copy()\n",
    "    # Assign custom residual variables\n",
    "    ds[\"B-MCO\"] = ds[\"B_NEC\"] - ds[\"B_NEC_CHAOS-Core\"]\n",
    "    ds[\"B-MCO-MLI\"] = ds[\"B-MCO\"] - ds[\"B_NEC_CHAOS-Static\"]\n",
    "    ds[\"B-MCO-MLI-MMA\"] = ds[\"B-MCO-MLI\"] - ds[\"B_NEC_CHAOS-MMA-Primary\"]\\\n",
    "                                          - ds[\"B_NEC_CHAOS-MMA-Secondary\"]\n",
    "    # Create a new DataArray to contain all these residual combinations\n",
    "    da = ds[\"B_NEC\"].copy()\n",
    "    da.name = \"B_residuals\"\n",
    "    # Expand to 3 dimensions for each residual combination to use\n",
    "    da = da.expand_dims({\"residuals\": 3})\n",
    "    da.coords[\"residuals\"] = [\"B-MCO\", \"B-MCO-MLI\", \"B-MCO-MLI-MMA\"]\n",
    "    # Assign the residual data to the DataArray\n",
    "    da = da.copy()\n",
    "    da.loc[{\"residuals\": \"B-MCO\"}] = ds[\"B-MCO\"]\n",
    "    da.loc[{\"residuals\": \"B-MCO-MLI\"}] = ds[\"B-MCO-MLI\"]\n",
    "    da.loc[{\"residuals\": \"B-MCO-MLI-MMA\"}] = ds[\"B-MCO-MLI-MMA\"]\n",
    "    # Assign the new DataArray to the original Dataset\n",
    "    ds[\"B_residuals\"] = da\n",
    "    return ds\n",
    "\n",
    "ds = assign_residuals(ds)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we create scatter plots to show the data in different views. We select only the first month - this reduces the crowding on the figures but shows that we should seek different methods to produce summary views of the data for longer time periods. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "facetgrid = (\n",
    "    ds.sel({\"Timestamp\": slice(\"2019-01-01\", \"2019-02-01\")})\n",
    "      .plot.scatter(x=\"QDLat\", y=\"B_residuals\", col=\"NEC\", row=\"residuals\",\n",
    "                    sharex=\"all\", sharey=\"row\", s=1, linewidths=0,\n",
    "                    hue=\"SunZenithAngle\", cmap=\"hot_r\"))\n",
    "for ax in facetgrid.axes.flat:\n",
    "    ax.set_facecolor(\"grey\")\n",
    "    ax.grid()\n",
    "    ax.set_ylim((-200, 200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    ds\n",
    "    .sel({\"Timestamp\": slice(\"2019-01-01\", \"2019-02-01\")})\n",
    "    .plot.scatter(x=\"Longitude\", y=\"Latitude\", hue=\"B-MCO-MLI-MMA\",\n",
    "                  s=1, linewidths=0, cmap=\"Spectral\", vmin=-50, vmax=50)\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    ds\n",
    "    .sel({\"Timestamp\": slice(\"2019-01-01\", \"2019-02-01\")})\n",
    "    .plot.scatter(x=\"MLT\", y=\"QDLat\", hue=\"B-MCO-MLI-MMA\",\n",
    "                  s=1, linewidths=0, cmap=\"Spectral\", vmin=-50, vmax=50)\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Statistical summaries of data over the full year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    ds\n",
    "       .groupby_bins(\"QDLat\", 90)\n",
    "       .std()[\"B-MCO-MLI-MMA\"]\n",
    "       .plot.line(x=\"QDLat_bins\")\n",
    ")\n",
    "plt.suptitle(\"Standard deviations\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows the large variability in the data over the polar regions"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 1200
  },
  "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
}