{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e932fe5a",
   "metadata": {},
   "source": [
    "# Swarm Measurements\n",
    "\n",
    "Following on from the previous notebook, *Making a main field model*, in this demo we will access Swarm data and make a model in a similar way using data from space instead of ground observatories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c2349e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import holoviews as hv\n",
    "import hvplot.xarray\n",
    "\n",
    "from viresclient import SwarmRequest\n",
    "from chaosmagpy.model_utils import design_gauss, synth_values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45cd9777",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Fetching data to use\n",
    "\n",
    "We fetch data as follows:\n",
    "- `MAGA_LR` product from Swarm: this is the magnetic low rate (1Hz) data from Swarm Alpha\n",
    "- the `B_NEC` variable is the field vector in the NEC (North, East, Centre) frame\n",
    "- `MLT`: Magnetic Local Time of each measurement\n",
    "- downsampled to one measurement per 10 seconds (`PT10S`)\n",
    "- retaining only nominal measurements (according to `Flags_B` and `Flags_F`)\n",
    "- during geomagnetically quiet times (Kp ≤ 2)\n",
    "- during January 2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311d6f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "request = SwarmRequest()\n",
    "request.set_collection(\"SW_OPER_MAGA_LR_1B\")\n",
    "request.set_products(\n",
    "    measurements=[\"B_NEC\"],\n",
    "    sampling_step=\"PT10S\",\n",
    "    auxiliaries=[\"MLT\"],\n",
    ")\n",
    "request.set_range_filter(\"Flags_B\", 0, 1)\n",
    "request.set_range_filter(\"Flags_F\", 0, 1)\n",
    "request.set_range_filter(\"Kp\", 0, 2)\n",
    "data = request.get_between(\n",
    "    start_time=dt.datetime(2020,1,1),\n",
    "    end_time=dt.datetime(2020,2,1)\n",
    ")\n",
    "ds = data.as_xarray()\n",
    "ds.attrs.pop(\"Sources\")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a71a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ds = ds.drop(\"Timestamp\").sel(NEC=\"C\")\n",
    "_ds.hvplot.scatter(\n",
    "    x=\"Longitude\", y=\"Latitude\", color=\"B_NEC\",\n",
    "    rasterize=True, colorbar=True, cmap=\"coolwarm\", clabel=\"Vertical (downwards) magnetic field (nT)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "626c8424",
   "metadata": {},
   "source": [
    "Above, we show the vertical (downwards) component of the magnetic field vector sampled across Earth. If you zoom in, you can see the lines curving as the orbits converge over the poles. As we have taken a whole month of data, the spatial coverage is quite dense."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271c9c6c",
   "metadata": {},
   "source": [
    "## Building a model\n",
    "\n",
    "Just as before, we perform a least-squares inversion to find Gauss coefficients.\n",
    "\n",
    "In this case, we can easily go to higher spherical harmonic degree because we have many more data and they are globally distributed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3350c987",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(ds, nmax=13):\n",
    "    \"\"\"Use the contents of dataset, ds, to build a SH model\"\"\"\n",
    "    # Get the positions and measurements from ds\n",
    "    radius = (ds[\"Radius\"]/1e3).values\n",
    "    theta = (90-ds[\"Latitude\"]).values\n",
    "    phi = (ds[\"Longitude\"]).values\n",
    "    B_radius = -ds[\"B_NEC\"].sel(NEC=\"C\").values\n",
    "    B_theta = -ds[\"B_NEC\"].sel(NEC=\"N\").values\n",
    "    B_phi = ds[\"B_NEC\"].sel(NEC=\"E\").values\n",
    "    # Use ChaosMagPy to build the design matrix, G\n",
    "    # https://chaosmagpy.readthedocs.io/en/master/functions/chaosmagpy.model_utils.design_gauss.html\n",
    "    G_radius, G_theta, G_phi = design_gauss(radius, theta, phi, nmax=nmax)\n",
    "    # Using only the radial component, perform the inversion:\n",
    "    G = G_radius\n",
    "    d = B_radius\n",
    "    m = np.linalg.inv(G.T @ G) @ (G.T @ d)\n",
    "    return m\n",
    "\n",
    "model_coeffs = build_model(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78397cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The first 10 coefficients\n",
    "model_coeffs[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea3db47",
   "metadata": {},
   "source": [
    "Let's sample our model and visualise it with contours:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "234dea9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_model_on_grid(m, radius=6371.200):\n",
    "    \"\"\"Evaluate Gauss coefficients over a regular grid over Earth\"\"\"\n",
    "    theta = np.arange(1, 180, 1)\n",
    "    phi = np.arange(-180, 180, 1)\n",
    "    theta, phi = np.meshgrid(theta, phi)\n",
    "    B_r, B_t, B_p = synth_values(m, radius, theta, phi)\n",
    "    ds_grid = xr.Dataset(\n",
    "        data_vars={\"B_NEC_model\": ((\"y\", \"x\", \"NEC\"), np.stack((-B_t, B_p, -B_r), axis=2))},\n",
    "        coords={\"Latitude\": ((\"y\", \"x\"), 90-theta), \"Longitude\": ((\"y\", \"x\"), phi), \"NEC\": np.array([\"N\", \"E\", \"C\"])}\n",
    "    )\n",
    "    return ds_grid\n",
    "\n",
    "# Evaluate at the mean radius of the satellite measurements\n",
    "mean_radius_km = float(ds[\"Radius\"].mean())/1e3\n",
    "ds_grid = sample_model_on_grid(model_coeffs, radius=mean_radius_km)\n",
    "# Generate contour plot of vertical component\n",
    "ds_grid.sel(NEC=\"C\").hvplot.contourf(\n",
    "    x=\"Longitude\", y=\"Latitude\", c=\"B_NEC_model\",\n",
    "    levels=30, coastline=True, projection=ccrs.PlateCarree(), global_extent=True,\n",
    "    clabel=\"Model: vertical (downwards) magnetic field (nT)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a69d4ae",
   "metadata": {},
   "source": [
    "Nice! It looks like the main field: a curved magnetic equator dropping southwards over South America (around the South Atlantic Anomaly where the field is weaker), the South magnetic pole occurs over the edge of Antarctica towards Australia, the North magnetic pole is more extended over Northern Canada and Siberia. Though it is not as accurate as published models!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "735b68f8",
   "metadata": {},
   "source": [
    "## Data-model residuals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e934098a",
   "metadata": {},
   "source": [
    "Now let's dig a little deeper, comparing the input data and the model...\n",
    "\n",
    "Let's plot the difference between the data and the model, the \"data-model residuals\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d045ca9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the model at the same points as the measurements\n",
    "def append_model_evaluations(ds, m):\n",
    "    ds = ds.copy()\n",
    "    B_r, B_t, B_p = synth_values(m, ds[\"Radius\"]/1e3, 90-ds[\"Latitude\"], ds[\"Longitude\"])\n",
    "    ds[\"B_NEC_model\"] = (\"Timestamp\", \"NEC\"), np.stack((-B_t, B_p, -B_r), axis=1)\n",
    "    ds[\"B_NEC_res_model\"] = ds[\"B_NEC\"] - ds[\"B_NEC_model\"]\n",
    "    return ds\n",
    "ds = append_model_evaluations(ds, model_coeffs)\n",
    "\n",
    "# Plot the \n",
    "_ds = ds.drop(\"Timestamp\").sel(NEC=\"C\")\n",
    "_ds.hvplot.scatter(\n",
    "    x=\"Longitude\", y=\"Latitude\", color=\"B_NEC_res_model\",\n",
    "    rasterize=True, colorbar=True, cmap=\"coolwarm\", clim=(-40, 40), clabel=\"Vertical (B_C) data-model residuals (nT)\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5813025d",
   "metadata": {},
   "source": [
    "(Above): The difference between the model and the data is quite small except over the poles... let's look at the other components now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b2958e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ds = ds.drop(\"Timestamp\").sel(NEC=\"E\")\n",
    "_ds.hvplot.scatter(\n",
    "    x=\"Longitude\", y=\"Latitude\", color=\"B_NEC_res_model\",\n",
    "    rasterize=True, colorbar=True, cmap=\"coolwarm\", clim=(-40, 40), clabel=\"Eastward (B_E) data-model residuals (nT)\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d855371",
   "metadata": {},
   "source": [
    "(Above): The disturbance over the poles is much stronger in the Eastward component... this is related to the geometry of the field-aligned currents producing a larger magnetic perturbation in that direction.\n",
    "\n",
    "Next let's plot against magnetic local time (MLT) instead, as the relevant coordinate system in which these currents are organised:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc15ff07-4ca3-4049-8056-debfd601790e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ds = ds.drop(\"Timestamp\").sel(NEC=\"E\")\n",
    "_ds.hvplot.scatter(\n",
    "    x=\"MLT\", y=\"Latitude\", color=\"B_NEC_res_model\",\n",
    "    rasterize=True, colorbar=True, cmap=\"coolwarm\", clim=(-40, 40), clabel=\"Eastward (B_E) data-model residuals (nT)\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe4c450",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.drop(\"Timestamp\").hvplot.scatter(x=\"Latitude\", y=\"B_NEC_res_model\", col=\"NEC\", rasterize=True, colorbar=False, cmap=\"bwr\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45bad123",
   "metadata": {},
   "source": [
    "(Above): Both the Northward (N) and Eastward (E) components are more disturbed over the poles than the downward (C) component. This disturbance is due to Field-Aligned Currents (FACs) and Auroral Electrojets (AEJs) which produce strong magnetic fields over the poles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "680ecdd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "(ds.groupby_bins(\"Latitude\", 90)\n",
    "   .apply(lambda x: x[\"B_NEC_res_model\"].std(axis=0))\n",
    "   .plot.line(x=\"Latitude_bins\", ax=ax)\n",
    ")\n",
    "ax.set_title(\"Standard deviations\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91840fc3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Exercises\n",
    "\n",
    "- Plot the power spectrum of our model\n",
    "- See what happens if you try push it to higher SH degree\n",
    "- Alter the inversion to use other components (the one here only used the vertical component)\n",
    "\n",
    "## Exploring further\n",
    "\n",
    "Take a look at [Swarm Notebooks](https://notebooks.vires.services/notebooks/04a1_geomag-models-vires) if you want more examples. There you will also find recipes for accessing other Swarm products."
   ]
  }
 ],
 "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",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}