{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9b9b7878",
   "metadata": {},
   "source": [
    "# Plotting\n",
    "\n",
    "`emsarray` can produce simple plots, to help you visualise your data quickly.\n",
    "Plots can be made from any supported format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8685add",
   "metadata": {},
   "outputs": [],
   "source": [
    "import emsarray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a627a012",
   "metadata": {},
   "source": [
    "Set up the environment..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7725caf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This makes the figures nice and big for this notebook\n",
    "from matplotlib import pyplot as plt\n",
    "plt.rcParams['figure.dpi'] = 100\n",
    "plt.rcParams['figure.figsize'] = (8, 5)\n",
    "plt.rcParams['savefig.dpi'] = 100\n",
    "\n",
    "# The coastlines used in the maps have some bad polygons,\n",
    "# which cause some warnings. Lets turn those off.\n",
    "import warnings\n",
    "warnings.filterwarnings(\n",
    "    \"ignore\", message=\"Shapefile shape has invalid polygon\",\n",
    "    category=UserWarning, module=\"shapefile\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405ffa77",
   "metadata": {},
   "source": [
    "A quick demonstration before we get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c42fcc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbr4 = emsarray.tutorial.open_dataset('gbr4')\n",
    "gbr4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3aada2",
   "metadata": {},
   "source": [
    "The `botz` variable indicates the depth of the ocean floor at a particular cell.\n",
    "\n",
    "The data will be plotted on a map, wih coastlines overlayed.\n",
    "A colour ramp will be applied and automatically scaled to the range of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aed8d2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbr4.ems.plot('botz', title='Depth of sea-bed')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a990294f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Selecting what to plot\n",
    "\n",
    "A plot has to be made of a single layer.\n",
    "For datasets with multiple depth layers and multiple time steps,\n",
    "start by selecting just one slice from each of these extra dimensions.\n",
    "\n",
    "The BRAN2020 example file contains multiple time steps and multiple depth layers.\n",
    "We will need to select just one of each before we can make a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09893dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "bran2020 = emsarray.tutorial.open_dataset('bran2020')\n",
    "bran2020"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2909e2b4",
   "metadata": {},
   "source": [
    "Using `dataset.isel(...)` we can trim the data so we have just the surface layer at the first timestep.\n",
    "After that we plot the `temp` variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fad1bd0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "bran2020.isel(Time=0, st_ocean=0).ems.plot('temp')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62fae437",
   "metadata": {},
   "source": [
    "Instead of selecting a subset of the data across the entire dataset,\n",
    "you can slice a single data array and plot that variable directly.\n",
    "To demonstrate, we can plot the last time step by selecting it on the data array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99cf0138",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = bran2020['temp'].isel(Time=-1, st_ocean=0)\n",
    "bran2020.ems.plot(temp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a477a8c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Plotting vectors\n",
    "\n",
    "You can plot vector data using two variables that make up the *u* and *v* components of the vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646a3830-79ca-40c6-ae6f-5a0664de2c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "fraser = emsarray.tutorial.open_dataset('fraser')\n",
    "fraser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec73c5d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select the surface currents at the first timestep\n",
    "fraser_0 = fraser.isel(time=0, k=-1)\n",
    "\n",
    "fraser.ems.plot(\n",
    "    vector=(fraser_0[\"u\"], fraser_0[\"v\"]),\n",
    "    title=\"Surface current velocity\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd885f9f",
   "metadata": {},
   "source": [
    "This can be combined with a scalar plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396ffade",
   "metadata": {},
   "outputs": [],
   "source": [
    "fraser.ems.plot(\n",
    "    scalar=fraser_0[\"eta\"],\n",
    "    vector=(fraser_0[\"u\"], fraser_0[\"v\"]),\n",
    "    title=\"Sea surface elevation and surface current velocty\",)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a40cbb72",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Plotting computed data\n",
    "\n",
    "Any data can be plotted on the model grid, as long as the data has the same shape.\n",
    "To demonstrate this, let us make a plot of surface temperature for our American friends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20c669dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "farhenheit_temp = fraser_0['temp'] * 1.8 + 32\n",
    "farhenheit_temp.attrs.update({\n",
    "    'long_name': 'Surface temperature (°F)',\n",
    "    'units': 'degrees Farhenheit'\n",
    "})\n",
    "fraser.ems.plot(farhenheit_temp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c06bd8",
   "metadata": {},
   "source": [
    "## Custom plots\n",
    "\n",
    "`dataset.ems.plot()` is useful for making quick plots of your data,\n",
    "but it does not create publication quality plots, and it does not expose any styling options.\n",
    "The following example shows how make a custom plot with your own styling, base layers, and titles.\n",
    "It uses a few tools provided by `emsarray` to extract the data and construct the polygons,\n",
    "but otherwise does all of the styling manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696f38c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cartopy.crs as ccrs\n",
    "import cartopy.io.img_tiles as cimgt\n",
    "import emsarray\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from emsarray.plot import polygon_to_patch, bounds_to_extent\n",
    "from emsarray.utils import datetime_from_np_time\n",
    "from matplotlib.collections import PatchCollection\n",
    "from shapely.geometry import box\n",
    "\n",
    "# What to plot\n",
    "dataset = emsarray.tutorial.open_dataset(\"fraser\")\n",
    "dataset = dataset.isel(time=0, k=-1)\n",
    "scalar = 'eta'\n",
    "vector = ('u', 'v')\n",
    "title = \"Sea surface height and surface current velocity\"\n",
    "datetime_stamp = datetime_from_np_time(dataset[\"time\"].values).strftime(\"%Y-%m-%d %H:%M\")\n",
    "\n",
    "# Constrain the view to this region\n",
    "view_area = box(152.3, -26.8, 154, -23.5)\n",
    "\n",
    "# What coordinate reference system the data are in\n",
    "data_crs = dataset.ems.data_crs\n",
    "\n",
    "# What coordinate reference system to plot the data in.\n",
    "# You can choose any projection that you want.\n",
    "view_crs = ccrs.NearsidePerspective(\n",
    "    central_longitude=view_area.centroid.x, central_latitude=view_area.centroid.y)\n",
    "\n",
    "# Set up a figure and some axes, then set view bounds\n",
    "figure = plt.figure(figsize=(8, 10), dpi=100)\n",
    "axes = plt.subplot(projection=view_crs)\n",
    "axes.set_title(title + \"\\n\" + datetime_stamp)\n",
    "axes.set_aspect(aspect='equal', adjustable='datalim')\n",
    "axes.set_extent(bounds_to_extent(view_area.bounds), crs=data_crs)\n",
    "\n",
    "# Add some grid lines\n",
    "gridlines = axes.gridlines(draw_labels=True, color=(0.1, 0.1, 0.1, 0.1))\n",
    "gridlines.top_labels = False\n",
    "gridlines.right_labels = False\n",
    "\n",
    "# Add a base image layer\n",
    "axes.add_image(cimgt.Stamen(style='terrain-background', desired_tile_form='RGBA'), 9)\n",
    "# axes.add_image(cimgt.Stamen(style='terrain-labels', desired_tile_form='RGBA'), 9)\n",
    "figure.text(\n",
    "    0.5, 0, \"Map tiles by Stamen Design, under CC BY 3.0. Data by OpenStreetMap, under ODbL.\",\n",
    "    va='bottom', ha='center', fontsize=8)\n",
    "\n",
    "# Plot the data\n",
    "patches = dataset.ems.make_patch_collection(scalar, edgecolor='face', cmap='plasma')\n",
    "axes.add_collection(patches)\n",
    "figure.colorbar(patches, ax=axes, location='right', label='metres')\n",
    "\n",
    "quiver = dataset.ems.make_quiver(axes, *vector)\n",
    "axes.add_collection(quiver)\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}