{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c5e01a29",
   "metadata": {},
   "source": [
    "# Clipping data\n",
    "\n",
    "`emsarray` provides utilities to clip data to some arbitrary geographic area.\n",
    "Any supported file format can be clipped using an identical API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6b77d6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import emsarray\n",
    "import shapely.geometry\n",
    "import tempfile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefa79aa",
   "metadata": {},
   "source": [
    "Set up the environment..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d63b45a5",
   "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['savefig.dpi'] = 100\n",
    "\n",
    "# The coastlines used in the maps have some bad polygons,\n",
    "# which causes 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": "8a1a0301",
   "metadata": {},
   "source": [
    "To demonstrate the data slicing, we need a custom plot function.\n",
    "Refer to the `clip.ipynb` notebook for further information about plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7426132b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from emsarray.plot import polygon_to_patch\n",
    "from matplotlib import patches, pyplot, collections\n",
    "\n",
    "def plot_clipped_on_grid(original, clipped, variable, clip_geometry):\n",
    "    \"\"\"Plot a variable from a clipped dataset\n",
    "    on top of the grid of the entire dataset,\n",
    "    plus the clip polygon\n",
    "    \"\"\"\n",
    "    figure = pyplot.figure()\n",
    "    clipped.ems.plot_on_figure(figure, variable)\n",
    "    axes = figure.axes[0]\n",
    "\n",
    "    grid_cells = original.ems.make_patch_collection(\n",
    "        edgecolor='grey', linewidth=0.5, facecolor='none', alpha=0.5, zorder=0)\n",
    "    axes.add_collection(grid_cells)\n",
    "\n",
    "    clip_patch = polygon_to_patch(\n",
    "        clip_geometry, edgecolor=\"red\", fill=False, linewidth=2, zorder=10)\n",
    "    axes.add_patch(clip_patch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e1cd2f2",
   "metadata": {},
   "source": [
    "First, open a dataset and define a clip region.\n",
    "A clip region can be any Shapely geometry.\n",
    "Typically this will be a polygon,\n",
    "but you can also use points, lines, multipolygons, and any other shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d92dd0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "gbr4 = emsarray.tutorial.open_dataset(\"gbr4\")\n",
    "whitsundays_clip = shapely.geometry.Polygon([\n",
    "    (147.391621, -20.346789), \n",
    "    (149.892177, -23.045809),\n",
    "    (153.440517, -20.680525),\n",
    "    (151.430263, -18.277805),\n",
    "    (147.391621, -20.346789),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4bacfb7",
   "metadata": {},
   "source": [
    "The dataset can be clipped to this region using `dataset.ems.clip(...)`.\n",
    "\n",
    "This operation needs a temporary directory to save intermediate files in to.\n",
    "We use `tempfile.TemporaryDirectory()` to manage this for us.\n",
    "\n",
    "Before the temporary directory is cleaned up we need to do one of two things.\n",
    "For small datasets or small clip regions we can call `dataset.load()`\n",
    "to load the entire clipped dataset in to memory.\n",
    "For larger datasets or clip regions that don't fit in to memory,\n",
    "we need to save dataset to a new file using `dataset.to_netcdf()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d71d986a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as temp_name:\n",
    "    whitsundays = gbr4.ems.clip(whitsundays_clip, temp_name)\n",
    "    whitsundays.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a6ae87",
   "metadata": {},
   "source": [
    "We can now plot this clipped dataset using th `plot_clipped_on_grid()` function we defined earlier.\n",
    "The coloured data regions come from the clipped dataset,\n",
    "while the grid lines are taken from the original dataset, to demonstrate what was removed.\n",
    "\n",
    "Any cell that intersects the clip region is selected.\n",
    "All neighbouring cells of these selected cells are added to the selected set, to form a buffer.\n",
    "Every other cell is masked with `nan`,\n",
    "and the dataset extents trimmed to the minimum bounding region of the clip polygon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77889bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_clipped_on_grid(gbr4, whitsundays.isel(time=0, k=-1), \"temp\", whitsundays_clip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2633734-99e9-4c5b-86ce-1378bced3e54",
   "metadata": {},
   "source": [
    "Any supported dataset format can be clipped, including datasets defined on an nstructured mesh. We can roughly cut out a region from the AUSTEn dataset around Tasmania:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac20a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "austen = emsarray.tutorial.open_dataset(\"austen\")\n",
    "tasmania_clip = shapely.geometry.Polygon([\n",
    "    (141.459, -40.780),\n",
    "    (142.954, -39.198),\n",
    "    (149.106, -39.095),\n",
    "    (150.864, -41.376),\n",
    "    (149.809, -44.621),\n",
    "    (144.843, -45.706),\n",
    "    (141.723, -43.389),\n",
    "    (141.372, -40.913),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ba3a4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as temp_name:\n",
    "    tasmania = austen.ems.clip(tasmania_clip, temp_name)\n",
    "    tasmania.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14d8595a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_clipped_on_grid(austen, tasmania, \"Mesh2_depth\", tasmania_clip)"
   ]
  }
 ],
 "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
}