{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f329e4f1-6dd9-4f2e-b5b3-0e8454832eee",
   "metadata": {},
   "source": [
    "# Dataset\n",
    "This notebook provides an example of the data utilised to generate a forecast using Aardvark Weather. We explore a single timeslice containing all the observations required to generate a forecast. This sample data is the output of the loader `WeatherDatasetE2E` in `../aardvark/loaders.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c17595b2-a624-440b-a2be-6af8d3802f87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bc8b869d-33d4-45e5-b3ef-15d4e521e3d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6de4a11f-4677-4552-8630-6b395c27012e",
   "metadata": {},
   "source": [
    "## Data\n",
    "Open a sample of data required to produce a forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0483c01d-2e08-469c-b914-42532cc3099b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/sample_data_final.pkl', 'rb') as fp:\n",
    "    data = pickle.load(fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83b7c966-fab5-47ce-a91c-d5419acaee40",
   "metadata": {},
   "source": [
    "Multiple different datasets are utilised as input to create a forecast, each with multiple channels including observations and metadata. Example channels for each of these are plotted below. The plot_channel variable in each cell can be adjusted to visualise different channels. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "693aa212-c3ab-4eda-b07a-a6b262da627d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up longitude and latitude for plotting\n",
    "lon = np.linspace(0,359,360)\n",
    "lat = np.linspace(90,-90,181)\n",
    "\n",
    "lon_1p5 = np.linspace(0,359,240)\n",
    "lat_1p5 = np.linspace(90,-90,121)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a2524f-108a-4ed7-bc89-2fb91638a7aa",
   "metadata": {},
   "source": [
    "## Satellite\n",
    "First visualise the satellite data from HIRS, AMSU-A, AMSU-B, IASI and ASCAT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1243f422-351a-458c-aaa9-4bf72a959678",
   "metadata": {},
   "source": [
    "### HIRS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d750dcc-cfef-4ca5-8f14-a7bae45980f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 11\n",
    "p = plt.contourf(\n",
    "    lon,\n",
    "    lat,\n",
    "    data[\"assimilation\"][\"hirs_current\"][0,...,plot_channel].cpu().T, \n",
    "    levels=100,\n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised radiance') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"HIRS channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c5e48e-3bab-461f-8425-9efa10c415e0",
   "metadata": {},
   "source": [
    "### AMSU-A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79deb284-7ed8-47d2-9006-a13a06d2b835",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 8\n",
    "p = plt.contourf(\n",
    "    lon, \n",
    "    lat[:-1],\n",
    "    data[\"assimilation\"][\"amsua_current\"][0,...,plot_channel].cpu(), \n",
    "    levels=100,\n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised radiance') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"AMSU-A channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b481a5c-97ad-4016-826d-93771531ad2d",
   "metadata": {},
   "source": [
    "### AMSU-B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "466434e7-7296-4454-b14c-6ab40bfbe0d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 10\n",
    "p = plt.contourf(\n",
    "    lon,\n",
    "    lat,\n",
    "    data[\"assimilation\"][\"amsub_current\"][0,...,plot_channel].T.cpu(), \n",
    "    levels=100,\n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised radiance') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"AMSU-A channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44de7c3f-6a25-47eb-bcfe-ab7bc82a1aac",
   "metadata": {},
   "source": [
    "### IASI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a5caf9-ecc4-403d-97ff-6c96dbc18f13",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 10\n",
    "p = plt.contourf(\n",
    "    lon,\n",
    "    lat,\n",
    "    data[\"assimilation\"][\"iasi_current\"][0,...,plot_channel].T.cpu(), \n",
    "    levels=100,\n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised radiance') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"IASI channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0ac772-4150-4bfd-96de-004ac7b8b6b4",
   "metadata": {},
   "source": [
    "### ASCAT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9981b9c6-2d88-47a1-b135-f70f1335df78",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 5\n",
    "p = plt.contourf(\n",
    "    lon,\n",
    "    lat,\n",
    "    data[\"assimilation\"][\"ascat_current\"][0,...,plot_channel].T.cpu(), \n",
    "    levels=100,\n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised radiance') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"ASCAT channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b6849d1-c6fd-413f-a72e-3de111228a78",
   "metadata": {},
   "source": [
    "## SYNOPS\n",
    "We next visualise the SYNOPS data from land stations, marine platforms and radiosonde profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c392b951-b446-4fde-aacb-02be7759bc97",
   "metadata": {},
   "source": [
    "### HadISD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611eed9e-8c4a-4ac8-82f4-9969a6ef8eb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plot_channel = 0\n",
    "p = plt.scatter(\n",
    "    data[\"assimilation\"][\"x_context_hadisd_current\"][plot_channel][0,0,:].cpu(),\n",
    "    data[\"assimilation\"][\"x_context_hadisd_current\"][plot_channel][0,1,:].cpu(),\n",
    "    c = data[\"assimilation\"][\"y_context_hadisd_current\"][plot_channel][0].T.cpu(), \n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised value') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"HadISD channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2ae8c40-28bd-49ad-8a2f-b9f7df706834",
   "metadata": {},
   "source": [
    "### ICOADS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e425818-5a30-4641-abc7-08cfd4c6b4fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_channel = 1\n",
    "fig = plt.figure()\n",
    "p = plt.scatter(\n",
    "    data[\"assimilation\"][\"icoads_x_current\"][0][0,:].cpu(),\n",
    "    data[\"assimilation\"][\"icoads_x_current\"][1][0,:].cpu(),\n",
    "    c = data[\"assimilation\"][\"icoads_current\"][0,plot_channel,:].T.cpu(), \n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised value') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"ICOADS channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14d103a9-053c-4576-b17f-35ebae3fa44a",
   "metadata": {},
   "source": [
    "### IGRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d5f5de-22ea-4d6e-a590-c1ce0a842a53",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_channel = 1\n",
    "fig = plt.figure()\n",
    "p = plt.scatter(\n",
    "    data[\"assimilation\"][\"igra_x_current\"][0][0,:].cpu(),\n",
    "    data[\"assimilation\"][\"igra_x_current\"][1][0,:].cpu(),\n",
    "    c = data[\"assimilation\"][\"igra_current\"][0,plot_channel,:].T.cpu(), \n",
    "    cmap=\"magma\")\n",
    "cbar = fig.colorbar(p)\n",
    "cbar.set_label('Normalised value') \n",
    "plt.xlabel(\"Longitude\")\n",
    "plt.ylabel(\"Latitude\")\n",
    "plt.title(f\"IGRA channel {plot_channel}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b893a77-4c30-4f18-a7fb-b007a676d849",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "521c93b6-babf-41fe-8dc6-c9a6eaef5b6f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.18"
  },
  "vscode": {
   "interpreter": {
    "hash": "ff379d8b57ea10bc9abbe40475203c855480f57b5b4821ba344a2e4e1a8a455a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}