{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo FACxTMS_2F (single satellite)\n",
    "\n",
    "> Authors: Ashley Smith\n",
    ">\n",
    "> Abstract: Access to the field aligned currents evaluated by the single satellite method (level 2 product). We show simple line plots of the time series over short periods (minutes), from both Swarm Alpha and Charlie. We also compare with the alternative method whereby the FACs are evaluated locally by computing them from the magnetic field data (`B_NEC` from `MAGx_LR_1B`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Documentation:\n",
    "\n",
    "- https://earth.esa.int/web/guest/missions/esa-eo-missions/swarm/data-handbook/level-2-product-definitions#FACxTMS_2F\n",
    "- https://earth.esa.int/documents/10174/1514862/Swarm_L2_FAC_single_product_description"
   ]
  },
  {
   "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 numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as mdates\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "\n",
    "request = SwarmRequest()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check what \"FAC\" data variables are available\n",
    "\n",
    "NB: these are the same as in the `FAC_TMS_2F` dual-satellite FAC product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.available_collections(\"FAC\", details=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.available_measurements(\"FAC\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting as a time series\n",
    "\n",
    "### Fetch one day from Swarm Alpha and Charlie\n",
    "\n",
    "Also fetch the quasidipole (QD) coordinates and Orbit Number at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request.set_collection(\"SW_OPER_FACATMS_2F\", \"SW_OPER_FACCTMS_2F\")\n",
    "request.set_products(\n",
    "    measurements=[\"FAC\", \"FAC_Error\", \n",
    "                  \"Flags\", \"Flags_F\", \"Flags_B\", \"Flags_q\"],\n",
    "    auxiliaries=[\"QDLat\", \"QDLon\", \"OrbitNumber\"],\n",
    ")\n",
    "data = request.get_between(\n",
    "    dt.datetime(2014,4,20),\n",
    "    dt.datetime(2014,4,21)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The source files of the original data are listed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data can be loaded as a pandas dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = data.as_dataframe()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively we can load the data as an xarray Dataset, though in the following examples we use the data via a pandas DataFrame instead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = data.as_xarray()\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on your application, you should probably do some filtering according to each of the flags. This can be done on the dataframe here, or beforehand on the server using [`request.set_range_filter()`](https://viresclient.readthedocs.io/en/latest/api.html#viresclient.SwarmRequest.set_range_filter). See https://earth.esa.int/documents/10174/1514862/Swarm_L2_FAC_single_product_description for more about the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the time series (FAC and FAC_Error for Alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(ncols=1, nrows=2, figsize=(15,5))\n",
    "# Select out the time series from Swarm Alpha\n",
    "dfA = df.where(df[\"Spacecraft\"] == \"A\").dropna()\n",
    "axes[0].plot(dfA.index, dfA[\"FAC\"])\n",
    "axes[1].plot(dfA.index, dfA[\"FAC_Error\"], color=\"orange\")\n",
    "axes[0].set_ylabel(\"FAC\\n[$\\mu A / m^2$]\");\n",
    "axes[1].set_ylabel(\"Error\\n[$\\mu A / m^2$]\");\n",
    "axes[1].set_xlabel(\"Timestamp\");\n",
    "date_format = mdates.DateFormatter('%Y-%m-%d\\n%H:%M')\n",
    "axes[1].xaxis.set_major_formatter(date_format)\n",
    "axes[0].set_ylim(-5, 5);\n",
    "axes[1].set_ylim(0, 1);\n",
    "axes[0].set_xticklabels([])\n",
    "fig.subplots_adjust(hspace=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot a subset of the time series (FAC from Alpha and Charlie)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def line_plot(fig, ax, df, varname=\"FAC\", spacecraft=\"A\", color=\"red\"):\n",
    "    \"\"\"Plot FAC as a line, given a dataframe\"\"\"\n",
    "    df = df.copy()\n",
    "    df = df.where(df[\"Spacecraft\"] == spacecraft).dropna()\n",
    "    ax.plot(df.index, df[varname], linewidth=1,\n",
    "            label=f\"{varname}$_{spacecraft}$\", color=color)\n",
    "    # Plot error range as filled area\n",
    "    if varname is \"FAC\":\n",
    "        ax.fill_between(df.index, \n",
    "                        df[\"FAC\"] - df[\"FAC_Error\"],\n",
    "                        df[\"FAC\"] + df[\"FAC_Error\"], color=\"grey\")\n",
    "    # Adjust limits and label formatting\n",
    "    datetime_format = \"%Y-%m-%d\\n%H:%M:%S\"\n",
    "    xlabel_format = mdates.DateFormatter(datetime_format)\n",
    "    ax.xaxis.set_major_formatter(xlabel_format)\n",
    "    ax.set_ylabel(\"[ $\\mu A / m^2$ ]\")\n",
    "    # Make y-axis symmetric about zero\n",
    "    ylim = max(abs(y) for y in ax.get_ylim())\n",
    "    ax.set_ylim((-ylim, ylim))\n",
    "    ax.legend()\n",
    "    ax.grid(True)\n",
    "    # Set up an extra xaxis at the top, to display Latitude\n",
    "    ax2 = ax.twiny()\n",
    "    ax2.set_xlim(ax.get_xlim())\n",
    "    ax2.set_xticks(ax.get_xticks())\n",
    "    # Identify closest times in dataframe to use for Latitude labels\n",
    "    # NB need to draw the figure now in order to get the xticklabels\n",
    "    #  https://stackoverflow.com/a/41124884\n",
    "    fig.canvas.draw()\n",
    "    # Extract times from the lower x axis\n",
    "    # Use them to find the nearest Lat values in the dataframe\n",
    "    xtick_times = [dt.datetime.strptime(ts.get_text(), datetime_format) for ts in ax.get_xticklabels()]\n",
    "    ilocs = [df.index.get_loc(t, method=\"nearest\") for t in xtick_times]\n",
    "    lats = df.iloc[ilocs][\"Latitude\"]\n",
    "    lat_labels = [\"{}°\".format(s) for s in np.round(lats.values, decimals=1)]\n",
    "    ax2.set_xticklabels(lat_labels)\n",
    "    ax2.set_xlabel(\"Latitude\")\n",
    "\n",
    "# Easy pandas-style slicing of the dataframe\n",
    "df_subset = df['2014-04-20T04:26:00':'2014-04-20T04:30:00']\n",
    "fig, axes = plt.subplots(nrows=2, figsize=(15, 5))\n",
    "line_plot(fig, axes[0], df_subset, spacecraft=\"A\", color=\"red\")\n",
    "line_plot(fig, axes[1], df_subset, spacecraft=\"C\", color=\"blue\")\n",
    "fig.subplots_adjust(hspace=0.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FAC estimates from (top) Swarm Alpha and (bottom) Swarm Charlie. The error estimate is shown as a thin grey area"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Also show satellite location on a map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def line_plot_figure(df, spacecraft=\"A\", color=\"red\"):\n",
    "    \"\"\"Generate a figure containing both line plot and maps\"\"\"\n",
    "    df = df.copy()\n",
    "    df = df.where(df[\"Spacecraft\"] == spacecraft).dropna()\n",
    "    # Set up figure geometry together with North/South maps\n",
    "    fig = plt.figure(figsize=(20, 5))\n",
    "    ax_lineplot = plt.subplot2grid((1, 5), (0, 0), colspan=3, fig=fig)\n",
    "    ax_N = plt.subplot2grid((1, 5), (0, 3), fig=fig,\n",
    "        projection=ccrs.Orthographic(\n",
    "            central_longitude=0.0, central_latitude=90.0\n",
    "        ))\n",
    "    ax_S = plt.subplot2grid((1, 5), (0, 4), fig=fig,\n",
    "        projection=ccrs.Orthographic(\n",
    "            central_longitude=0.0, central_latitude=-90.0\n",
    "        ))\n",
    "    for _ax in (ax_N, ax_S):\n",
    "        _ax.set_global()\n",
    "        _ax.coastlines(color=\"grey\")\n",
    "        _ax.add_feature(cfeature.LAND)\n",
    "        _ax.add_feature(cfeature.OCEAN)\n",
    "        _ax.plot(df[\"Longitude\"], df[\"Latitude\"], transform=ccrs.PlateCarree(),\n",
    "                 linewidth=4, color=color)\n",
    "    # Draw the line plot as before\n",
    "    line_plot(fig, ax_lineplot, df, spacecraft=spacecraft, color=color)\n",
    "\n",
    "line_plot_figure(df_subset, spacecraft=\"A\", color=\"red\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing to values calculated locally from the B_NEC data (experimental)\n",
    "\n",
    "ref https://nbviewer.jupyter.org/github/smithara/viresclient_examples/blob/master/swarmpyfac_dev.ipynb \n",
    "\n",
    "NB this is currently fixed to the time window and satellite combination (A and C) set above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from swarmpyfac.fac import single_sat_fac\n",
    "\n",
    "def _fetch_data(start='2014-04-20T04:26:00',\n",
    "                end='2014-04-20T04:30:02', spacecraft=\"A\"):\n",
    "    \"\"\"Fetch B_NEC data and combined geomagnetic model for a given spacecraft\"\"\"\n",
    "    request = SwarmRequest()\n",
    "    request.set_collection(f\"SW_OPER_MAG{spacecraft}_LR_1B\")\n",
    "    request.set_products(\n",
    "        measurements=[\"B_NEC\"],\n",
    "        models=['Model = \"MCO_SHA_2C\" + \"MLI_SHA_2C\" + \"MMA_SHA_2C-Primary\" + \"MMA_SHA_2C-Secondary\"'])\n",
    "    data = request.get_between(start, end,\n",
    "                             asynchronous=False, show_progress=False)\n",
    "    return data.as_xarray()\n",
    "\n",
    "def _get_input_data_from_ds(ds):\n",
    "    \"\"\"Extract the \"input_data\" required from the dataset\"\"\"\n",
    "    # convert time to unix time seconds\n",
    "    ## if on a dataframe, df:\n",
    "    ## time = np.array(df.index.astype(np.int64) / 10**9)\n",
    "    # on xarray.Dataset:\n",
    "    time = ds['Timestamp'].data.astype(np.int64) / 10**9\n",
    "    theta = ds['Latitude'].data\n",
    "    phi = ds['Longitude'].data\n",
    "    r = ds['Radius'].data\n",
    "    # equivalent to swarmpyfac.utils.pack_3d:\n",
    "    positions = np.stack((theta, phi, r), axis=1)\n",
    "    B_model = ds['B_NEC_Model'].data\n",
    "    B_res = ds['B_NEC'].data - B_model\n",
    "    return {'time': time, 'positions':positions,\n",
    "            'B_res': B_res, 'B_model': B_model}\n",
    "\n",
    "def _append_fac(ds=None):\n",
    "    \"\"\"Append FAC calculations to a dataset\"\"\"\n",
    "    input_data = _get_input_data_from_ds(ds)\n",
    "    output = single_sat_fac(**input_data)\n",
    "    # outputs like these should probably be turned into a dict so that they can be identified\n",
    "    irc = output[2]\n",
    "    fac = output[3]\n",
    "    time = output[0]\n",
    "    # Append the new data to the dataset\n",
    "    # https://xarray.pydata.org/en/stable/data-structures.html#dictionary-like-methods\n",
    "    # Note that there must now be a new offset time coordinate\n",
    "    ds.coords['Timestamp_2'] = pd.to_datetime(time, unit='s')\n",
    "    ds[f'FAC_calculated'] = (('Timestamp_2',), fac)\n",
    "    ds[f'IRC_calculated'] = (('Timestamp_2',), irc)\n",
    "    return ds\n",
    "\n",
    "def append_FAC_calculated_locally(df):\n",
    "    \"\"\"Use the functions above to evaluate FAC\n",
    "    \n",
    "    NB currently depends on the fixed time and spacecraft selection as before\n",
    "\n",
    "    \"\"\"\n",
    "    # Create xarray datasets containing the FAC estimates from each of A and C\n",
    "    ds_A = _append_fac(_fetch_data('2014-04-20T04:26:00',\n",
    "                                 '2014-04-20T04:30:02', spacecraft=\"A\"))\n",
    "    ds_C = _append_fac(_fetch_data('2014-04-20T04:26:00',\n",
    "                                 '2014-04-20T04:30:02', spacecraft=\"C\"))\n",
    "    # Transform them into a concatenated dataframe\n",
    "    df_A = pd.DataFrame(ds_A[\"FAC_calculated\"].to_pandas(), columns=[\"FAC_calc\"])\n",
    "    df_C = pd.DataFrame(ds_C[\"FAC_calculated\"].to_pandas(), columns=[\"FAC_calc\"])\n",
    "    df_calc = pd.concat((df_A, df_C))\n",
    "    df_calc.index.name = \"\"\n",
    "    # Append them to the existing dataframe\n",
    "    df = df.copy()\n",
    "    df[\"FAC_new\"] = df_calc[\"FAC_calc\"]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate FACs locally and append to the dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_subset = append_FAC_calculated_locally(df_subset)\n",
    "df_subset[\"FAC_diff\"] = df_subset[\"FAC\"] - df_subset[\"FAC_new\"]\n",
    "df_subset.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the difference between the FAC sourced from the product, and the one evaluated locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=4, figsize=(15, 10))\n",
    "line_plot(fig, axes[0], df_subset, varname=\"FAC\", spacecraft=\"A\", color=\"red\")\n",
    "line_plot(fig, axes[0], df_subset, varname=\"FAC_new\", spacecraft=\"A\", color=\"blue\")\n",
    "line_plot(fig, axes[1], df_subset, varname=\"FAC_diff\", spacecraft=\"A\", color=\"black\")\n",
    "line_plot(fig, axes[2], df_subset, varname=\"FAC\", spacecraft=\"C\", color=\"red\")\n",
    "line_plot(fig, axes[2], df_subset, varname=\"FAC_new\", spacecraft=\"C\", color=\"blue\")\n",
    "line_plot(fig, axes[3], df_subset, varname=\"FAC_diff\", spacecraft=\"C\", color=\"black\")\n",
    "fig.subplots_adjust(hspace=0.8)"
   ]
  }
 ],
 "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
}