{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo MAGxLR_1B (magnetic field 1Hz)\n",
    "\n",
    "> Authors: Ashley Smith\n",
    ">\n",
    "> Abstract: Access to the low rate (1Hz) magnetic data (level 1b product), together with geomagnetic model evaluations (level 2 products)."
   ]
  },
  {
   "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 datetime as dt\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "request = SwarmRequest()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MAGX_LR_1B product information\n",
    "\n",
    "This is one of the main products from Swarm - the 1Hz measurements of the magnetic field vector (`B_NEC`) and total intensity (`F`). These are derived from the Vector Field Magnetometer (VFM) and Absolute Scalar Magnetomer (ASM).\n",
    "\n",
    "Documentation:\n",
    "- https://earth.esa.int/web/guest/missions/esa-eo-missions/swarm/data-handbook/level-1b-product-definitions#MAGX_LR_1B_Product\n",
    "\n",
    "Measurements are available through VirES as part of collections with names containing `MAGx_LR`, for each Swarm spacecraft:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.available_collections(\"MAG\", details=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The measurements can be used together with geomagnetic model evaluations as shall be shown below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check what \"MAG\" data variables are available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.available_measurements(\"MAG\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the names of available models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.available_models(details=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetch one hour of MAG data and models, at 10-second sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.set_collection(\"SW_OPER_MAGA_LR_1B\")\n",
    "request.set_products(\n",
    "    measurements=[\"F\", \"B_NEC\"],\n",
    "    models=[\"CHAOS-Core\", \"MCO_SHA_2D\"],\n",
    "    sampling_step=\"PT10S\"\n",
    ")\n",
    "data = request.get_between(\n",
    "    # 2014-01-01 00:00:00\n",
    "    start_time = dt.datetime(2014,1,1, 0),\n",
    "    # 2014-01-01 01:00:00\n",
    "    end_time = dt.datetime(2014,1,1, 1)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### See a list of the input files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transfer data to a pandas dataframe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = data.as_dataframe()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `expand=True` to extract vectors (B_NEC...) as separate columns (..._N, ..._E, ..._C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = data.as_dataframe(expand=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ... or to an xarray Dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = data.as_xarray()\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.Sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instead, fetch the residuals directly\n",
    "\n",
    "Adding `residuals=True` to `.set_products()` will instead directly evaluate and return all data-model residuals"
   ]
  },
  {
   "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=[\"F\", \"B_NEC\"],\n",
    "    models=[\"CHAOS-Core\", \"MCO_SHA_2D\"],\n",
    "    residuals=True,\n",
    "    sampling_step=\"PT10S\"\n",
    ")\n",
    "data = request.get_between(\n",
    "    start_time = dt.datetime(2014,1,1, 0),\n",
    "    end_time = dt.datetime(2014,1,1, 1)\n",
    ")\n",
    "df = data.as_dataframe(expand=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the scalar residuals for each model\n",
    "\n",
    "### ... using the pandas method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = df.plot(\n",
    "    y=[\"F_res_CHAOS-Core\", \"F_res_MCO_SHA_2D\"],\n",
    "    figsize=(15,5),\n",
    "    grid=True\n",
    ")\n",
    "ax.set_xlabel(\"Timestamp\")\n",
    "ax.set_ylabel(\"[nT]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ... using matplotlib interface (Matlab-style)\n",
    "\n",
    "NB: we are doing `plt.plot(x, y)` with `x` as `df.index` (the time-based index of df), and `y` as `df[\"..\"]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,5))\n",
    "plt.plot(\n",
    "    df.index,\n",
    "    df[\"F_res_CHAOS-Core\"],\n",
    "    label=\"F_res_CHAOS-Core\"\n",
    ")\n",
    "plt.plot(\n",
    "    df.index,\n",
    "    df[\"F_res_MCO_SHA_2D\"],\n",
    "    label=\"F_res_MCO_SHA_2D\"\n",
    ")\n",
    "plt.xlabel(\"Timestamp\")\n",
    "plt.ylabel(\"[nT]\")\n",
    "plt.grid()\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ... using matplotlib interface (Object Oriented style)\n",
    "\n",
    "This is the recommended route for making more complicated figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15,5))\n",
    "ax.plot(\n",
    "    df.index,\n",
    "    df[\"F_res_CHAOS-Core\"],\n",
    "    label=\"F_res_CHAOS-Core\"\n",
    ")\n",
    "ax.plot(\n",
    "    df.index,\n",
    "    df[\"F_res_MCO_SHA_2D\"],\n",
    "    label=\"F_res_MCO_SHA_2D\"\n",
    ")\n",
    "ax.set_xlabel(\"Timestamp\")\n",
    "ax.set_ylabel(\"[nT]\")\n",
    "ax.grid()\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the vector components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(15,10), sharex=True)\n",
    "for component, ax in zip(\"NEC\", axes):\n",
    "    for model_name in (\"CHAOS-Core\", \"MCO_SHA_2D\"):\n",
    "        ax.plot(\n",
    "            df.index,\n",
    "            df[f\"B_NEC_res_{model_name}_{component}\"],\n",
    "            label=model_name\n",
    "        )\n",
    "    ax.set_ylabel(f\"{component}\\n[nT]\")\n",
    "    ax.legend()\n",
    "axes[0].set_title(\"Residuals to models (NEC components)\")\n",
    "axes[2].set_xlabel(\"Timestamp\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Similar plotting, using the data via xarray instead\n",
    "\n",
    "xarray provides a more sophisticated data structure that is more suitable for the complex vector data we are accessing, together with nice stuff like unit and other metadata support. Unfortunately due to the extra complexity, this can make it difficult to use right away."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = data.as_xarray()\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(15,10), sharex=True)\n",
    "for i, ax in enumerate(axes):\n",
    "    for model_name in (\"CHAOS-Core\", \"MCO_SHA_2D\"):\n",
    "        ax.plot(\n",
    "            ds[\"Timestamp\"],\n",
    "            ds[f\"B_NEC_res_{model_name}\"][:, i],\n",
    "            label=model_name\n",
    "        )\n",
    "    ax.set_ylabel(\"NEC\"[i] + \" [nT]\")\n",
    "    ax.legend()\n",
    "axes[0].set_title(\"Residuals to models (NEC components)\")\n",
    "axes[2].set_xlabel(\"Timestamp\");\n",
    "# automatic unit labels will be possible in v0.5.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that xarray also allows convenient direct plotting like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"B_NEC_res_CHAOS-Core\"].plot.line(x=\"Timestamp\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Access multiple MAG datasets simultaneously\n",
    "\n",
    "It is possible to fetch data from multiple collections simultaneously. Here we fetch the measurements from Swarm Alpha and Bravo. In the returned data, you can differentiate between them using the \"Spacecraft\" column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_MAGA_LR_1B\", \"SW_OPER_MAGC_LR_1B\")\n",
    "request.set_products(\n",
    "    measurements=[\"F\", \"B_NEC\"],\n",
    "    models=[\"CHAOS-Core\",],\n",
    "    residuals=True,\n",
    "    sampling_step=\"PT10S\"\n",
    ")\n",
    "data = request.get_between(\n",
    "    start_time = dt.datetime(2014,1,1, 0),\n",
    "    end_time = dt.datetime(2014,1,1, 1)\n",
    ")\n",
    "df = data.as_dataframe(expand=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[df[\"Spacecraft\"] == \"A\"].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[df[\"Spacecraft\"] == \"C\"].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ... or using xarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = data.as_xarray()\n",
    "ds.where(ds[\"Spacecraft\"] == \"A\", drop=True)"
   ]
  }
 ],
 "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
}