Image_Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc3fd7b9-4c3e-42c1-871e-b6c2455d4a8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import IFrame # Display YouTube videos"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c822981-ecc1-4cbd-9e1f-d8e86e015b03",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ccffcc; padding: 10px;\">\n",
    "    <h1> Tutorial X </h1>\n",
    "    <h2> Image Segmentation with U-Net and fastai </h2>\n",
    "</div>\n",
    "\n",
    "# Overview\n",
    "\n",
    "This Jupyter notebook demonstrates how artificial neural networks (ANNs) can be applied to image segmentation problems. Segmentation in this context refers to the task of assigning discrete labels to individual pixels or regions of an image. We can use segmentation models to identify and locate features of interest within an image. This notebook contains a simple application to self-driving cars, where we train a segmentation model to identify important features in dashcam footage, as well as a more complicated example, based on the work of [Coney et al. (2023)](https://doi.org/10.1002/qj.4592), identifying and characterising trapped lee waves over the UK.\n",
    "\n",
    "## Recommended reading\n",
    "\n",
    "* [Fastai: A Layered API for Deep Learning](https://doi.org/10.3390/info11020108)\n",
    "* [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://doi.org/10.1007/978-3-319-24574-4_28)\n",
    "* [Identifying and characterising trapped lee waves using deep learning techniques](https://doi.org/10.1002/qj.4592)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5036e3b3-9002-4c96-8398-76b6178d1df4",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "<div style=\"background-color: #e6ccff; padding: 10px;\">\n",
    "\n",
    "<h1> Machine Learning Theory </h1>\n",
    "\n",
    "# Image Segmentation\n",
    "\n",
    "## The problem\n",
    "\n",
    "Image segmentation models are designed to tackle the problem of partitioning an image into meaningful segments or regions, each corresponding to different objects or parts of objects within the image. This process is crucial in various applications such as medical imaging, where it helps in identifying and isolating different anatomical structures (e.g. organs or tumours), or in autonomous driving, where it can aid in recognising and distinguishing between pedestrians, vehicles, and road signs. By accurately segmenting images, these models enable more precise analysis and interpretation, facilitating tasks like object detection, scene understanding, and image editing. Essentially, image segmentation transforms raw visual data into structured information, making it easier for machines to understand and interact with the visual world. More recently, segmentation models are being applied to weather and climate forecasting applications, where their ability to identify structures in image data makes them ideally suited.\n",
    "\n",
    "## Popular models for image segmentation\n",
    "\n",
    "* U-Net: Its architecture has become a standard in medical image segmentation due to its ability to perform well with limited training data and its precise localization capabilities.\n",
    "* Mask R-CNN: This model is highly significant for instance segmentation, as it not only detects objects but also provides pixel-level masks, making it versatile for various applications, including autonomous driving and video analysis.\n",
    "* DeepLab: Known for its high accuracy in semantic segmentation, DeepLab’s use of atrous convolution allows it to capture multi-scale context, making it a powerful tool for tasks requiring detailed scene understanding."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0077789e-4d1f-4de3-92a9-d8527c060816",
   "metadata": {},
   "source": [
    "## The U-Net model architecture\n",
    "\n",
    "The [U-Net](https://doi.org/10.1007/978-3-319-24574-4_28) model architecture is a type of convolutional neural network (CNN) originally designed for biomedical image segmentation. Introduced by Olaf Ronneberger, Philipp Fischer, and Thomas Brox in 2015, U-Net is known for its distinctive U-shaped structure. This architecture consists of a contracting path (the encoder) to capture context and a symmetric expanding path (the decoder) that enables precise localization. The contracting path follows the typical architecture of a convolutional network, with repeated application of convolutions, each followed by a rectified linear unit (ReLU) and a max-pooling operation. The expanding path, on the other hand, involves upsampling the feature maps and performing convolutions, which helps in reconstructing the image with high resolution. U-Net's ability to work with very few training images and its efficient use of data augmentation make it particularly effective for tasks where annotated data is scarce.\n",
    "\n",
    "![Schematic of U Net](https://rmets.onlinelibrary.wiley.com/cms/asset/63d9263f-f5c7-48dc-8ba6-60f19fb6e5a7/qj4592-fig-0004-m.jpg)\n",
    "\n",
    "The video in the cell below gives a 10-minute introduction to the U-Net."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24326866-49a9-4140-bc3c-98dc6f51c164",
   "metadata": {
    "id": "OfUQlcc_b6Ky"
   },
   "outputs": [],
   "source": [
    "IFrame(\"https://www.youtube.com/embed/NhdzGfB1q74?si=p8ti5ydxXvqJuABi\",\"560\", \"315\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3d263b-f2ff-4f3b-8080-f5fcb13b2e06",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "# Python\n",
    "\n",
    "## [PyTorch](https://pytorch.org/)\n",
    "\n",
    "PyTorch is an open-source machine learning library developed by Facebook's AI Research lab. It is widely used for applications such as natural language processing and computer vision. PyTorch is known for its flexibility and ease of use, particularly due to its dynamic computation graph, which allows for more intuitive model building and debugging. However, it is considered quite low-level compared to some other frameworks (e.g. [Keras](https://keras.io/)), meaning that defining complex models like a U-Net can require a significant amount of verbose code. This verbosity can make the development process more cumbersome, especially for those who are new to deep learning.\n",
    "\n",
    "## [fastai](https://www.fast.ai/)\n",
    "\n",
    "fastai is a high-level library built on top of PyTorch that simplifies the process of training deep learning models. It provides a range of pre-built functions and classes that allow users to leverage the powerful capabilities of PyTorch without needing to write extensive amounts of code. With fastai, you can define and train complex models, such as U-Nets, in just a few lines of code. This makes it an excellent choice for both beginners and experienced practitioners who want to quickly prototype and experiment with different models while still benefitting from the flexibility and performance of PyTorch under the hood.\n",
    "\n",
    "## Further reading\n",
    "\n",
    "If you want to run this notebook locally or on a remote service:\n",
    "\n",
    "* [running Jupyter notebooks](https://jupyter.readthedocs.io/en/latest/running.html)\n",
    "* [installing the required Python environments](https://github.com/cemac/LIFD_ENV_ML_NOTEBOOKS/blob/main/howtorun.md)\n",
    "* [running the Jupyter notebooks locally](https://github.com/cemac/LIFD_ENV_ML_NOTEBOOKS/blob/main/jupyter_notebooks.md)\n",
    "\n",
    "</div>\n",
    "\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe41b425-f3c5-4a7f-8324-5353e9f54c4c",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #ffffcc; padding: 10px;\">\n",
    "    \n",
    "<h1> Requirements </h1>\n",
    "\n",
    "This notebook should run with the following requirements satisfied.\n",
    "\n",
    "<h2> Python Packages: </h2>\n",
    "\n",
    "* fastai\n",
    "* pytorch\n",
    "* numpy\n",
    "* xarray\n",
    "* dask\n",
    "* netCDF4\n",
    "* bottleneck\n",
    "* matplotlib\n",
    "* cartopy\n",
    "* notebook\n",
    "\n",
    "<h2> Data Requirements</h2>\n",
    "\n",
    "This notebook refers to some external datasets and learner objects which are downloaded via Python scripts within the notebook.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "518cca5f-bae7-4c95-b6a4-5c220d73b990",
   "metadata": {},
   "source": [
    "**Contents:**\n",
    "\n",
    "1. [Overview and machine-learning theory](#Overview)\n",
    "2. [Application to self-driving cars](#Application-to-self-driving-cars)\n",
    "3. [Application to detection of lee waves](#Application-to-detection-of-lee-waves)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19aa0f56-e4a6-44d7-a60f-568e8fbaea68",
   "metadata": {},
   "source": [
    "<div style=\"background-color: #cce5ff; padding: 10px;\">\n",
    "\n",
    "## Import modules\n",
    "\n",
    "These are all the modules needed during this tutorial.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adbe44a7-dd7f-4b36-9253-f9befffa56b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision.all import *\n",
    "from matplotlib.colors import ListedColormap\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import urllib.request\n",
    "import os\n",
    "import zipfile\n",
    "import pickle\n",
    "import fastai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdfdfb8e-2e36-46f1-83fb-0383ab48d48d",
   "metadata": {},
   "source": [
    "### Note on CUDA\n",
    "If you have a GPU then you should enable CUDA by commenting out the cell below. Otherwise, all code will run on the CPU by default (much slower)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20c5fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fastai.torch_core.default_device(use=False) # comment out if you have a GPU!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6621326-b510-4d83-9b54-199463f15270",
   "metadata": {},
   "source": [
    "### Note for Windows users\n",
    "\n",
    "The fastai library was designed with Linux filesystems in mind, so can raise certain errors on Windows machines. The following cell is a workaround to allow the rest of the notebook to run on Windows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a6803ef-8a37-48d4-b985-32e9d4fd9f1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.name == 'nt':\n",
    "    import pathlib\n",
    "    temp = pathlib.PosixPath\n",
    "    pathlib.PosixPath = pathlib.WindowsPath"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee58789",
   "metadata": {},
   "source": [
    "### Note for Colab users\n",
    "\n",
    "Google Colab doesn't come with cartopy installed by default. Uncomment the cell below to rectify this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7523518",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install cartopy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f5b1d9-bef4-423f-aa3f-38069379ee92",
   "metadata": {},
   "source": [
    "# Application to self-driving cars\n",
    "\n",
    "This example is a quick introduction to the `unet_learner` object in fastai and is based on the official package documentation.\n",
    "\n",
    "We will start by downloading a small version of the [CamVid](http://mi.eng.cam.ac.uk/research/projects/VideoRec/CamVid/) dataset. The Cambridge-driving Labeled Video Database (CamVid) is the first collection of videos with object class semantic labels, complete with metadata. The database provides ground truth labels that associate each pixel with one of 32 semantic classes. The image below uses colour to show how pixels are labelled according to the object present.\n",
    "\n",
    "![CamVid image](http://mi.eng.cam.ac.uk/research/projects/VideoRec/CamVid/pr/DBOverview1_1_huff_0000964.jpg)\n",
    "\n",
    "fastai includes many helper functions, such as `untar_data`, which simplify downloading of example datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2ce07e-e731-41d5-9fc9-59b0fcfce06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download subset of the CamVid segmentation dataset\n",
    "path = untar_data(URLs.CAMVID_TINY)\n",
    "path.ls()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1e4e669-8cf6-459b-a631-db1ea2562029",
   "metadata": {},
   "source": [
    "The above cell has printed the names of directories where the images and label masks are stored. The label assigned to each pixel is stored as an integer. Class labels corresponding to each integer are stored in `codes.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f35da26c-6fd6-4002-8c0b-ea76f08f7d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import class labels\n",
    "codes = np.loadtxt(path/'codes.txt', dtype=str)\n",
    "len(codes), codes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "733ac0bb-86c4-42fc-8c01-9e6f40ff7eed",
   "metadata": {},
   "source": [
    "We can see that there are 32 unique classes with corresponding labels printed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c252e758-a921-4c85-a45e-6a0038505e64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get filenames of input images\n",
    "fnames = get_image_files(path/'images')\n",
    "fnames[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c1ef751-ca11-444e-ae45-f9aaf5386e4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at example of label file\n",
    "(path/'labels').ls()[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a051e1c6-3ea9-411c-b217-309b8d90a183",
   "metadata": {},
   "source": [
    "Note that the image file containing the label mask has the same filename as the original image but with `_P` appended just before the `.png` file extension.\n",
    "\n",
    "We can now write a function called `label_func` to associate each image file with its corresponding label mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecdf0dda-e688-4a1b-9a3d-6e4fd2386f85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to locate file containing class labels\n",
    "def label_func(fn):\n",
    "    return path/'labels'/f'{fn.stem}_P{fn.suffix}'\n",
    "label_func(fnames[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74a17fa3-d3e7-45af-9035-426f94a5e379",
   "metadata": {},
   "source": [
    "The `DataBlock` object in fastai is a powerful and flexible tool designed to simplify the process of creating datasets for machine learning. It provides a high-level API that allows you to define the structure of your data pipeline in a clear and concise manner. With `DataBlock`, you can specify how to get your data, how to split it into training and validation sets, how to label it, and how to apply transformations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3861f044-64b7-4e52-a5c9-c83d6d01c6fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build fastai DataBlock\n",
    "camvid = DataBlock(\n",
    "    blocks=(ImageBlock, MaskBlock(codes)), # (input, output)\n",
    "    get_items=get_image_files, # function for retrieving input files\n",
    "    get_y=label_func, # function to locate label mask\n",
    "    splitter=RandomSplitter(), # how to perform train/validation split\n",
    "    batch_tfms=aug_transforms(size=(120,160)) # data augmentation transforms including output size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4985eea-d2dd-45f0-9e3d-8c3b023ebd12",
   "metadata": {},
   "source": [
    "Now that we have a `DataBlock`, we can create a `DataLoaders` object. This is similar to a `DataLoader` in PyTorch, but more high-level in that it manages more aspects of the training process. As well as our `DataBlock`, we need to specify the directory where our training data are located and the batch size to use in training.\n",
    "\n",
    "We are choosing a batch size of eight to ensure that VRAM usage doesn't exceed 4 GB. Depending on the size of your GPU memory, you may wish to increase the batch size to speed up training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2dfa56-bc1d-43f2-838b-fe23a046d7e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a DataLoaders object from the DataBlock (not the same as DataLoader in PyTorch)\n",
    "dls = camvid.dataloaders(path/'images', path=path, bs=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6538dc4-92d6-4459-98a1-05763f0948f3",
   "metadata": {},
   "source": [
    "To verify that our `DataLoaders` object is working correctly, we can display some images from a sample batch with their class labels overlaid as colours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29649ffb-218c-42fe-92ee-5bff244a8ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show segmented images from sample batch\n",
    "dls.show_batch(max_n=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfa8857d-caa7-41b5-8761-e00c950ad744",
   "metadata": {},
   "source": [
    "In the fastai library, a learner object is a central component designed to streamline the process of training and evaluating models. It encapsulates the model, the data, and the training loop, providing a high-level API that simplifies complex tasks. With a learner, you can easily fine-tune hyperparameters, apply callbacks, and leverage built-in functionalities like learning rate scheduling and mixed-precision training.\n",
    "\n",
    "For our purposes here, we will instantiate a `unet_learner` object. This will automatically set up the neural network for the dimensions of our problem, including the number of unique pixel classes. In addition to automating the training loop, the `unet_learner` uses a pre-trained neural network as its encoder model. Here we specify ResNet-34 as the encoder. This is a residual neural network with 34 layers and has been pre-trained on the ImageNet database of more than one million images!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cea4e56-18b4-44b9-a166-d3e65d42d2ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download fastai U-Net learner with ResNet-34 as the encoder\n",
    "learn = unet_learner(dls, resnet34)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72265339-7453-4ef5-8777-4733c4b017c9",
   "metadata": {},
   "source": [
    "fastai encourages users to leverage the concept of transfer learning, rather than training new models from scratch. In transfer learning, we typically download a model (here ResNet-34) that has been pre-trained on a large generic dataset to perform a standard task, e.g. image recognition. This pre-trained model, sometimes called a foundation model, contains a great deal of prior information for recognising features in images. For a specific task, such as identifying objects in dashcam footage, we need only fine-tune this foundation model (iteratively update the weights and biases) for a few epochs on our training data.\n",
    "\n",
    "Depending on your hardware, running the next cell may take a few minutes, so be warned! The next few lines of code are therefore commented out by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b317a6c5-b4cb-404b-aa5c-26984c5d9787",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fine-tune the model until the validation loss stops improving\n",
    "# learn.fine_tune(100, cbs=EarlyStoppingCallback(monitor='valid_loss', patience=5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bff628-b615-4409-960c-e7eb7f4b8209",
   "metadata": {},
   "source": [
    "After few epochs of fine tuning, our validation loss has stopped improving and the model is trained. We can now plot our model predictions superimposed on the original images, and compare with the true label masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e97127f-6de7-473d-8dfb-dc12e16a1709",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare model output with ground truth\n",
    "# learn.show_results(max_n=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd22cc4-8600-495c-9983-20cd3c718894",
   "metadata": {},
   "source": [
    "fastai includes helper functions such as `plot_top_losses`, which shows which examples the model had the hardest time classifying. Looking at the highest losses can be a good way to find outliers and errors in the training data (unlikely here, as this dataset is very clean)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47008642-8b5f-4d36-aad0-d29da0eaca00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at the images where the validation loss was highest\n",
    "# interp = SegmentationInterpretation.from_learner(learn)\n",
    "# interp.plot_top_losses(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a5c1c32-a758-42f5-8a48-04eaf14a1a7d",
   "metadata": {},
   "source": [
    "In this example, we have seen how, in only a few lines of code, we can train a segmentation model with good performance and with a relatively short training time. The model predictions, although not perfect, are remarkable given the small amount of training data and the high number of object classes (32).\n",
    "\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "410dbdf1-271b-4ba9-947d-7b5ec01b599f",
   "metadata": {},
   "source": [
    "# Application to detection of lee waves\n",
    "\n",
    "This example is based on the work of [Coney et al. (2023)](https://doi.org/10.1002/qj.4592), who used neural networks to identify and characterise trapped lee waves over the UK. A lee wave is type of gravity wave created when air in the atmosphere flows over mountainous terrain.\n",
    "\n",
    "## Gravity waves\n",
    "\n",
    "A gravity wave is a vertical wave, for example a ripple, in the atmosphere. Gravity waves can be formed when air is forced upwards by topography (e.g. wind blowing over a mountain). This creates turbulence that can be felt throughout the column of air above a mountain. Gravity waves are of interest for improving our understanding and forecasting capability, e.g. for aviation. If you'd like to learn more, NOAA have a useful information page all about gravity waves in the atmosphere [here](https://www.weather.gov/source/zhu/ZHU_Training_Page/Miscellaneous/gravity_wave/gravity_wave.html).\n",
    "\n",
    "![diagram of gravity waves](https://www.weather.gov/source/zhu/ZHU_Training_Page/Miscellaneous/gravity_wave/radarscope2.png)\n",
    "\n",
    "(taken from https://www.weather.gov/source/zhu/ZHU_Training_Page/Miscellaneous/gravity_wave/gravity_wave.html)\n",
    "\n",
    "## Lee waves\n",
    "\n",
    "Lee waves can be observed by eye as you get clouds forming on the crest of the wave, e.g. when you look up and see stripes of clouds or lenticular clouds like the image seen below, where a mountain has forced a wave in the air to form. These can be spotted in photos and satellite images. For more information, the Met Office have a basic overview [here](https://www.metoffice.gov.uk/weather/learn-about/weather/types-of-weather/wind/lee-waves).\n",
    "\n",
    "![Lenticular cloud over mountains image](https://www.metoffice.gov.uk/binaries/content/gallery/metofficegovuk/images/weather/learn-about/weather/lenticular-cloud.jpg)\n",
    "\n",
    "(taken from https://www.metoffice.gov.uk/weather/learn-about/weather/types-of-weather/wind/lee-waves)\n",
    "\n",
    "## NWP data\n",
    "\n",
    "Lee waves can be identified in Numerical Weather Prediction (NWP) model output in a range of fields, such as vertical wind velocity just above topography. Below is an image of model output where lee waves are resolved, showing a characteristic stripey vertical velocity pattern. These patterns are easily picked up by eye, but not so easily detected automatically. To detect these patterns, typically spectral analysis is employed using idealised representations of waves.\n",
    "\n",
    "![Example UKV data showing stripey lee waves in the verticle velocity output](https://rmets.onlinelibrary.wiley.com/cms/asset/10a1023d-3e98-4f26-9100-224ac84ea3d1/qj4592-fig-0001-m.jpg)\n",
    "\n",
    "(Figure 1 from [Coney et al. (2023)](https://doi.org/10.1002/qj.4592))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb73a37-3729-4d03-bcaf-2c9849d4d230",
   "metadata": {},
   "source": [
    "## Downloading the training data\n",
    "\n",
    "To start with, we need to download and extract the training data using the following two Python scripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55add8f9-681c-42ef-82db-2f231e3eefee",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://zenodo.org/records/10230764/files/data.zip\"\n",
    "filename = \"data.zip\"\n",
    "\n",
    "if not os.path.isfile(filename):\n",
    "    # If the file doesn't exist, download it using urllib\n",
    "    urllib.request.urlretrieve(url, filename)\n",
    "    print(\"File downloaded successfully.\")\n",
    "else:\n",
    "    print(\"File already exists.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188752ae-a959-46d6-a2ee-f8d318badad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the 'data' directory exists\n",
    "if not os.path.isdir('data'):\n",
    "    print(\"Directory 'data' does not exist. Creating it now...\")\n",
    "    \n",
    "    # Create the 'data' directory\n",
    "    os.mkdir('data')\n",
    "    \n",
    "    print(\"Extracting contents of 'data.zip' into current directory...\")\n",
    "    \n",
    "    # Extract the contents of 'data.zip' into the current directory\n",
    "    with zipfile.ZipFile('data.zip', 'r') as zip_ref:\n",
    "        # Extract only the 'data' directory from within the zip file\n",
    "        for member in zip_ref.namelist():\n",
    "            if member.startswith('data/'):\n",
    "                zip_ref.extract(member, '.')\n",
    "    \n",
    "    print(\"Extraction complete.\")\n",
    "else:\n",
    "    print(\"Directory 'data' already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61fe7f5c-19ab-4b71-9e8a-8526cc7cfbb8",
   "metadata": {},
   "source": [
    "Once you have the data already downloaded, re-running the above scripts won't do anything."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c38ce0-1ab6-4319-a59b-a208cd1f8c6b",
   "metadata": {},
   "source": [
    "## Training the segmentation model\n",
    "\n",
    "Setting our root path to the directory containing the training data will save us some typing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f841ec-3861-4309-9f29-d0c48833e260",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set root path to directory containing training data\n",
    "root = Path('data/train')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff1727f7-b011-4032-a73b-77eaf8969494",
   "metadata": {},
   "source": [
    "Unlike in the previous example, where the class label definitions were provided in a file that came with the data, this time we need to define the class labels ourselves. We do this by creating a Python dictionary called `codes`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902a85c7-6207-433a-9ee9-05757854adc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define binary class labels\n",
    "codes = {0:'no wave', 255:'lee wave'}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "774ca332-aca4-493e-a393-85ccb1383d87",
   "metadata": {},
   "source": [
    "Each pixel in our segmentation mask is therefore encoded as an 8-bit integer, with black pixels indicating a lee wave is present.\n",
    "\n",
    "Again, we need to define a function to match input images to their respective segmentation mask (take a peek inside the data directory to verify that this function will do what we intend...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0718ed23-866a-468b-a442-a241d5654543",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to retrieve label mask for a given input file\n",
    "def label_func(fn): \n",
    "    string = str(fn.stem)[:49] + 'mask.png' # mask files have .png suffix\n",
    "    return root/'masks_png'/string"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d651c950-cac8-40d0-9de0-d808a820608c",
   "metadata": {},
   "source": [
    "The NWP data that we will be using as input to our network is stored in NetCDF format, as is common for weather and climate datasets. fastai does not understand this file format natively, so we therefore need to define a function `open_xarray` to read in the data. The function uses the xarray library to open a NetCDF file and returns a NumPy array of the vertical velocities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e25fa9-b0c0-44df-84f1-16a3beb83f72",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to extract vertical velocity array from NetCDF file\n",
    "def open_xarray(fname):\n",
    "    x = xr.open_dataarray(fname)\n",
    "    array = x.values # return values as NumPy array\n",
    "    return array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb789260-a488-4de9-9c2b-74c3ed563a58",
   "metadata": {},
   "source": [
    "One of fastai's more powerful features is the way that it automates the process of data augmentation. Data augmentation is a technique used to increase the diversity of data available for training machine learning models by applying various transformations, such as rotations or flips, to existing data, thereby improving the model's performance and robustness.\n",
    "\n",
    "Here we apply several transformations to our data:\n",
    "\n",
    "* z-score normalisation centres and scales the pixel values about zero. This improves performance of the optimisation algorithms used for training the model.\n",
    "* random flipping\n",
    "* random zooming in\n",
    "* random rotation\n",
    "\n",
    "Augmenting the data using these transformations effectively increases the size of our training dataset and greatly reduces the chance of overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f4b9dd2-6dad-4b3d-b2f3-8cac506c8897",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data augmentation transformations to apply on GPU\n",
    "tfms = [\n",
    "    Normalize.from_stats([0,0,0], [1,1,1]), # mean zero and std dev. one for all channels\n",
    "    Flip(), # random flip images with probability 0.5\n",
    "    Zoom(max_zoom=20, p=0.5), # apply up to a 20x zoom with probability 0.5\n",
    "    Rotate(max_deg=360, p=0.9) # apply a random rotation with probability 0.9\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d16f827-2a78-4593-be97-c00c6b1a450e",
   "metadata": {},
   "source": [
    "We now have everything we need to build a fastai `DataBlock` object. Note that we are using the `open_xarray` and `label_func` functions that we just defined, as well as our dictionary `codes`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2713a38e-b4d1-4d7a-a2bd-29d158cb36bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build fastai DataBlock\n",
    "waves_ds = DataBlock(\n",
    "    blocks=(ImageBlock, MaskBlock(codes)), # (input, output)\n",
    "    get_items=get_files, # function for retrieving input files (not images this time!)\n",
    "    get_x=open_xarray, # function to extract input array from NetCDF file\n",
    "    get_y=label_func, # function to locate label mask\n",
    "    splitter=RandomSplitter(), # how to perform train/validation split\n",
    "    batch_tfms=tfms, # data augmentation transforms to be applied\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6953070c-f2f7-40fb-bdc4-0e90e0a8d2eb",
   "metadata": {},
   "source": [
    "A `DataLoaders` object can now be created as before. This time we set the batch to two: these images are higher resolution, so take up more VRAM per image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f394df45-1777-4848-a0d3-69776f28fa57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a DataLoaders object from the DataBlock\n",
    "dls = waves_ds.dataloaders(root/'vertical_velocities', path=root, bs=2) # batch size of two to conserve GPU memory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fffd4d1f-b35e-4e71-9e9b-78189019cf39",
   "metadata": {},
   "source": [
    "With the `show_batch` function we can check that the `DataLoaders` is working correctly. Note that some of the images in the training set do not contain lee waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "992aedfe-006d-4118-bf94-76de70f355d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show segmented images from sample batch (batch size is just two)\n",
    "dls.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e2e853-683f-4083-b0f1-afdd1f33687a",
   "metadata": {},
   "source": [
    "The models in this example will take much longer to train that the model in the first example. We will therefore be saving trained models to disc. The following Python script creates a directory to store the trained models, if such a directory doesn't yet exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97856f24-8300-44d4-95d1-a0a4f523bfec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the 'models_out' directory exists\n",
    "if not os.path.isdir('models_out'):\n",
    "    print(\"Directory 'models_out' does not exist. Creating it now...\")\n",
    "    \n",
    "    # Create the 'models_out' directory\n",
    "    os.mkdir('models_out')\n",
    "else:\n",
    "    print(\"Directory 'models_out' already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49f2cd69-5460-47ef-b3cc-bf469a7af062",
   "metadata": {},
   "source": [
    "We will use the same U-Net / ResNet-34 architecture as before. During training we will also monitor the [$F_1$ score](https://en.wikipedia.org/wiki/F-score) on the validation set. The $F_1$ score is the harmonic mean of the [precision and recall](https://en.wikipedia.org/wiki/Precision_and_recall) of a binary classifier. Values range from zero to one, with one indicating a perfect score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3105b42e-9e0c-4984-9e0f-046e81963fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download fastai U-Net learner with ResNet-34 as the encoder and compute F1 score\n",
    "learn2 = unet_learner(dls, resnet34, metrics=DiceMulti)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5224c05-63c6-4987-82eb-26d5b915046e",
   "metadata": {},
   "source": [
    "The learner object is now ready for training. Note that, unless you are using a very powerful GPU, the following code may take tens of minutes and possibly hours to run! The code is therefore commented out by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af04de1-d3e5-46b5-8e99-9a7d0d226a99",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Fine-tune the model until the validation loss stops improving\n",
    "# learn2.fine_tune(100, cbs=EarlyStoppingCallback(monitor='valid_loss', patience=5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5972258a-6cd4-4917-9cc5-63f8ff547e71",
   "metadata": {},
   "source": [
    "Each fastai learner stores a filesystem path as one of its attributes. We must first change this path to our `models_out` directory before exporting `learn2` as a pickle file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a0bf0a1-8214-4c85-8ff7-4c12db4ecd1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export trained model to 'models_out' directory\n",
    "# learn2.path = Path('models_out')\n",
    "# learn2.export('segmodel.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fea9536-6c92-4588-a068-1b50eb803373",
   "metadata": {},
   "source": [
    "A quick comparison of the model predictions with ground truth should indicate whether training has been successful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a33419-ad0c-4df5-b82a-59a777c98a7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare model output with ground truth\n",
    "# learn2.show_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd4f1ce4-6038-4b71-88e7-7a74d38dc007",
   "metadata": {},
   "source": [
    "## Training alternative model heads\n",
    "\n",
    "Using the segmentation model, we can identify where lee waves are present from the vertical velocity field. However, we might also be interested in the physical characteristics of those waves. For example:\n",
    "\n",
    "* amplitude\n",
    "* wavelength\n",
    "* orientation\n",
    "\n",
    "We could start from scratch and train a new model to predict each quantity, but a more efficient approach is to use transfer learning. Our segmentation model has already been trained to identify features related to lee waves. It therefore makes a perfect starting point to develop further models to predict wave characteristics. We can do this by training what are called alternative model heads, where head refers to the last few layers of the neural network. In general, the earlier layers of the network serve to extract relevant features from the image, while the head uses these features to predict a quantity of interest. New models can be cheaply obtained by simply training new heads on an old model.\n",
    "\n",
    "Unfortunately, the NWP data that we used to train the segmentation model does not include measurements of any of the wave characteristics we are interested in. To get around this problem, [Coney et al. (2023)](https://doi.org/10.1002/qj.4592) generated their own synthetic lee wave data using Leif Denby's [synthetic-gravity-waves](https://doi.org/10.5281/zenodo.7576811) package for Python. The following two Python scripts download and extract these synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73069842-9c08-459e-8990-f5a8e8f88398",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://huggingface.co/datasets/CEMAC/synthetic_lee_waves/resolve/main/synthetic_data.zip\"\n",
    "filename = \"synthetic_data.zip\"\n",
    "\n",
    "# Check if the file exists\n",
    "if not os.path.isfile(filename):\n",
    "    print(f\"File '{filename}' does not exist. Downloading it now...\")\n",
    "    \n",
    "    # Download the file using urllib\n",
    "    urllib.request.urlretrieve(url, filename)\n",
    "    \n",
    "    print(\"File downloaded successfully.\")\n",
    "else:\n",
    "    print(f\"File '{filename}' already exists.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0fe1b09-8eed-4e13-8573-1f966c339172",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the 'var_amp_synthetic' directory exists\n",
    "if not os.path.isdir('var_amp_synthetic'):\n",
    "    print(\"Directory 'var_amp_synthetic' does not exist. Creating it now...\")\n",
    "    \n",
    "    # Create the 'var_amp_synthetic' directory\n",
    "    os.mkdir('var_amp_synthetic')\n",
    "    \n",
    "    print(\"Extracting contents of 'synthetic_data.zip' into 'var_amp_synthetic' directory...\")\n",
    "    \n",
    "    # Extract only the 'var_amp_synthetic' contents from the zip file\n",
    "    with zipfile.ZipFile('synthetic_data.zip', 'r') as zip_ref:\n",
    "        for member in zip_ref.namelist():\n",
    "            if member.startswith('var_amp_synthetic/'):\n",
    "                zip_ref.extract(member, '.')\n",
    "    \n",
    "    print(\"Extraction complete.\")\n",
    "else:\n",
    "    print(\"Directory 'var_amp_synthetic' already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2cff614-f29d-41f1-9eb1-d1c2e06bf1ed",
   "metadata": {},
   "source": [
    "First, let's rebase our root path to the synthetic training data directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29c402b4-039b-4409-8395-66efc0b77e0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set root path to directory containing training data\n",
    "root = Path('var_amp_synthetic/train')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cbf4bd-ae6a-48d8-81b8-6829b865b8b3",
   "metadata": {},
   "source": [
    "Although we will not be using label masks for training our alternative model heads, we still need the labelling function to be defined in our global namespace. This is a peculiarity of the fastai library and is necessary for it to function properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86f5d816-b34a-4471-8e1a-2b1d7d41088f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to retrieve label mask for a given input file\n",
    "def label_func(fn): \n",
    "    string = str(fn.stem)[:49] + 'mask.png'\n",
    "    return root/'masks_png'/string"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56b53c05-9b9b-4059-824a-43d6ff0f67e1",
   "metadata": {},
   "source": [
    "The function to read input files is the same as before. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3749383f-497f-4f54-9f96-4f59f198f1f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to extract vertical velocity array from NetCDF file\n",
    "def open_xarray(fname):\n",
    "    x = xr.open_dataarray(fname)\n",
    "    array = x.values\n",
    "    return array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f123ba27-ee1b-4e44-8929-f1fee9745172",
   "metadata": {},
   "source": [
    "The synthetic lee wave data are stored in NumPy format. We therefore need to define a new function to read these files and return them in the format expected by our model. We also add some Gaussian noise to the data each time they are read in, as a form of data augmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09d2afb-0dfd-4e53-8e8d-e87457c4eee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to read in NumPy array file\n",
    "def open_np(fname):\n",
    "    x = np.load(fname)\n",
    "    noise = np.random.normal(size=(512, 512))\n",
    "    x = x + threshold*noise # add Gaussian noise\n",
    "    x2 = np.array([x, x, x]) # copy data to three channels for input into ResNet-34\n",
    "    return torch.Tensor(x2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c3a5a6f-eccd-4aae-9643-0dbfa445f91b",
   "metadata": {},
   "source": [
    "The files containing the different wave characteristics are stored in appropriately named directories. We need to define a separate retrieval function for each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f48baa-3296-49b1-bada-54941c07ad8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to retrieve corresponding wavelength labels\n",
    "def label_func_wl(fn):\n",
    "    string = str(fn.stem)[:49] + '.npy'\n",
    "    lbl = np.load(root/'wavelength'/string).astype('float')\n",
    "    return lbl/1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b94595-d4eb-4586-a717-85211d4e0809",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to retrieve corresponding amplitude labels\n",
    "def label_func_amp(fn):\n",
    "    string = str(fn.stem)[:49] + '.npy'\n",
    "    lbl = np.load(root/'amplitude'/string).astype('float')\n",
    "    return lbl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87debfe0-7548-4f47-b49c-b311dac641b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to retrieve corresponding orientation labels\n",
    "def label_func_or(fn):\n",
    "    string = str(fn.stem)[:49] + '.npy'\n",
    "    lbl = np.load(root/'orientation'/string).astype('float')\n",
    "    lbl_rad = lbl*np.pi/180 # convert degrees to radians\n",
    "    return np.array([np.sin(lbl_rad),np.cos(lbl_rad)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883ec9f5-dbe8-4da5-b787-fb360a0daff5",
   "metadata": {},
   "source": [
    "The procedure for training an alternative model head is sufficiently complicated to warrant its own function. To summarise briefly what this `train` function does:\n",
    "\n",
    "1. Our pre-trained segmentation model learner is loaded from disc.\n",
    "2. The `model` attribute is extracted (this contains the neural network itself).\n",
    "3. The orientation model will return sines and cosines of the orientation angle, so the dimensions of the output layer can remained unchanged (for segmentation the network was returning two probabilities).\n",
    "4. The amplitude and wavelength models will return a single number, so we need to redefine the last layers of the network to return one output.\n",
    "5. We are now predicting real numbers, not probabilities, so the loss function is changed to Mean Squared Error (MSE).\n",
    "6. A new `DataLoaders` is constructed with the modified model and loss function.\n",
    "7. Different learning rates are used for different layers.\n",
    "8. Weights and biases are frozen in the earlier layers of the model, so that only the head is updated during training.\n",
    "9. The trained model is saved to disc.\n",
    "\n",
    "We will call this function three times, once for each wave characteristic. Note that the threshold argument determines how much Gaussian noise is added to the inputs (this is a hyperparameter to tune).\n",
    "\n",
    "Note that these models take even longer to train than the segmentation model, so be warned! The offending lines of code are commented out by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67e4a18d-c05b-4788-8392-d275033d6249",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to train alternative model to predict wave characteristic\n",
    "def train(waves, characteristic, threshold, epochs=100):\n",
    "    dls = waves.dataloaders(root/'data', path=root, bs=2)\n",
    "    learn3 = load_learner('models_out/segmodel.pkl')\n",
    "    model = learn3.model # extract segmentation model from learner object\n",
    "    if characteristic != 'orientation':\n",
    "        model.layers[-2] = nn.Sequential(\n",
    "            torch.nn.Conv2d(99, 50, kernel_size=(1, 1), stride=(1, 1)),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Conv2d(50, 1, kernel_size=(1, 1), stride=(1, 1))\n",
    "        ) # redefine last layer of model head to give one-dimensional output for amplitude or wavelength\n",
    "    loss_func = MSELossFlat() # MSE loss function for real output\n",
    "    learn3 = Learner(dls, model, loss_func=loss_func) # build new learner object\n",
    "    base_lr = 1e-4\n",
    "    print('lr', base_lr)\n",
    "    lr_mult = 10 # multiplier for differential learning rates across layers\n",
    "    learn3.unfreeze()\n",
    "    learn3.freeze_to(-3) # only train model head (freeze parameters in earlier layers)\n",
    "    learn3.fit_one_cycle(epochs, slice(base_lr/lr_mult, base_lr), cbs=EarlyStoppingCallback(monitor='valid_loss', patience=5))\n",
    "    learn3.path = Path('models_out')\n",
    "    learn3.export(characteristic + '_' + str(threshold) + '.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b0d620-eee5-4450-9752-f9872872f9fe",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train amplitude model\n",
    "threshold = 0.0625\n",
    "waves = DataBlock(\n",
    "    blocks=(DataBlock, DataBlock), # use DataBlocks for arrays\n",
    "    get_items=get_files, # function for retrieving input files\n",
    "    get_x=open_np, # function to read NumPy array input file\n",
    "    get_y=label_func_amp, # function to locate amplitude labels\n",
    "    splitter=RandomSplitter(), # how to perform train/validation split\n",
    "    batch_tfms=[Normalize.from_stats(*imagenet_stats)], # normalize using mean and std devs from ImageNet dataset used to train ResNet-34\n",
    ")\n",
    "# train(waves, 'amplitude', threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1df7288f-443e-4e04-a21b-138ef213c20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train wavelength model\n",
    "threshold = 0.125\n",
    "waves = DataBlock(\n",
    "    blocks=(DataBlock, DataBlock), # use DataBlocks for arrays\n",
    "    get_items=get_files, # function for retrieving input files\n",
    "    get_x=open_np, # function to read NumPy array input file\n",
    "    get_y=label_func_wl, # function to locate wavelength labels\n",
    "    splitter=RandomSplitter(), # how to perform train/validation split\n",
    "    batch_tfms=[Normalize.from_stats(*imagenet_stats)], # normalize using mean and std devs from ImageNet dataset used to train ResNet-34\n",
    ")\n",
    "# train(waves, 'wavelength', threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a83fb27-b96f-4489-a76b-d2cb798b2798",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train orientation model\n",
    "threshold = 0.25\n",
    "waves = DataBlock(\n",
    "    blocks=(DataBlock, DataBlock), # use DataBlocks for arrays\n",
    "    get_items=get_files, # function for retrieving input files\n",
    "    get_x=open_np, # function to read NumPy array input file\n",
    "    get_y=label_func_or, # function to locate orientation labels\n",
    "    splitter=RandomSplitter(), # how to perform train/validation split\n",
    "    batch_tfms=[Normalize.from_stats(*imagenet_stats)], # normalize using mean and std devs from ImageNet dataset used to train ResNet-34\n",
    ")\n",
    "# train(waves, 'orientation', threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b675c557-26a4-4edf-9521-7ab63d32a0c3",
   "metadata": {},
   "source": [
    "## Plotting the results\n",
    "\n",
    "We have trained our U-Net segmentation model to identify trapped lee waves, as well as three alternative heads to predict wave characteristics. It is now time to visualise the results and see how our model performs on test data.\n",
    "\n",
    "Since training these models is very time-consuming, it is recommended to download the pre-trained models using the Python script below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9276e9ac-31a8-4502-a261-c846b3eabb6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output directory if it doesn't exist\n",
    "os.makedirs('models_out', exist_ok=True)\n",
    "\n",
    "urls = [\n",
    "    \"https://huggingface.co/CEMAC/LeeWaveNet/resolve/main/segmodel.pkl\",\n",
    "    \"https://huggingface.co/CEMAC/LeeWaveNet/resolve/main/amplitude_0.0625.pkl\",\n",
    "    \"https://huggingface.co/CEMAC/LeeWaveNet/resolve/main/wavelength_0.125.pkl\",\n",
    "    \"https://huggingface.co/CEMAC/LeeWaveNet/resolve/main/orientation_0.25.pkl\"\n",
    "]\n",
    "\n",
    "for url in urls:\n",
    "    filename = os.path.basename(url)\n",
    "    filepath = os.path.join('models_out', filename)\n",
    "\n",
    "    if not os.path.isfile(filepath):\n",
    "        # If the file doesn't exist, download it\n",
    "        urllib.request.urlretrieve(url, filepath)\n",
    "        print(f\"File {filename} downloaded successfully to models_out.\")\n",
    "    else:\n",
    "        print(f\"File {filename} already exists in models_out.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dac0838d-b579-4a06-b32d-b7bb09e8e9ad",
   "metadata": {},
   "source": [
    "It will be easier if we keep all our trained models in one place. Let's define a function to read them from disc and store them in a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7da00898-f33b-429f-b4b4-ec2b94c202a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to load the trained models\n",
    "def load_models():\n",
    "    learn2 = load_learner('models_out/segmodel.pkl')\n",
    "    wavelength_model = load_learner('models_out/wavelength_0.125.pkl')\n",
    "    orientation_model = load_learner('models_out/orientation_0.25.pkl')\n",
    "    amplitude_model = load_learner('models_out/amplitude_0.0625.pkl')\n",
    "    models_dict = {\n",
    "        'segmentation': learn2,\n",
    "        'wavelength': wavelength_model,\n",
    "        'orientation': orientation_model,\n",
    "        'amplitude': amplitude_model\n",
    "    }\n",
    "    return models_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f2aead5-1307-4621-be9a-0ebe86a16fd2",
   "metadata": {},
   "source": [
    "Making predictions from the trained models is slightly tricky. Due to fastai peculiarities, the vertical velocity input data need to be supplied in slightly different formats to the various models. We are storing the prediction output in `DataArray` objects using the xarray library, and we are returning these together in a `DataSet`. Adding the coordinates as metadata will make plotting easier later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e251848d-ec34-4e09-8f65-dc709b974817",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to make predictions with trained models\n",
    "def predict(models_dict, ds, xcoord='projection_x_coordinate', ycoord='projection_y_coordinate', mask_nonwaves=True):\n",
    "    arr = ds['upward_air_velocity'].values # extract input array of air velocities\n",
    "    \n",
    "    segmentation = models_dict['segmentation'].predict(arr)[0].numpy()\n",
    "    ds['segmentation'] = ((ycoord, xcoord), segmentation)\n",
    "    \n",
    "    wavelength = models_dict['wavelength'].predict(torch.Tensor([arr, arr, arr]))[0][0].numpy()\n",
    "    ds['wavelength'] = ((ycoord, xcoord), wavelength)\n",
    "    \n",
    "    orient = models_dict['orientation'].predict(torch.Tensor([arr, arr, arr]))[0]\n",
    "    orient = 180/np.pi * np.arctan(orient[0]/orient[1]) # convert sines and cosines into an angle in degrees\n",
    "    ds['orientation'] = ((ycoord, xcoord), orient)\n",
    "    \n",
    "    amplitude =  models_dict['amplitude'].predict(torch.Tensor(np.array([arr, arr, arr])))[0][0]\n",
    "    ds['amplitude'] = ((ycoord, xcoord), amplitude)\n",
    "    \n",
    "    if mask_nonwaves:\n",
    "        for char in ['amplitude','orientation','wavelength']:\n",
    "            ds[char] = ds[char].where(ds['segmentation']==1) # remove characteristic values where no waves\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c940abdb",
   "metadata": {},
   "source": [
    "The cell below defines some custom matplotlib colour maps that we will use for plotting our results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20aa9e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define colour map for vertical velocities\n",
    "vv_cmap = ListedColormap([(0/255,56/255,116/255),(0/255,101/255,206/255),(0/255,128/255,189/255),\n",
    "                         (0/255,151/255,154/255),(0/255,174/255,119/255),(0/255,197/255,85/255),\n",
    "                         (0/255,220/255,50/255),(0/255,244/255,16/255),(255/255,253/255,255/255),\n",
    "                         (238/255,212/255,0/255),(225/255,182/255,0/255),(213/255,151/255,0/255),\n",
    "                         (201/255,121/255,0/255),(189/255,91/255,0/255),(177/255,60/255,0/255),\n",
    "                         (165/255,30/255,0/255),(153/255,0/255,0/255),])\n",
    "\n",
    "# Define colour map for wave amplitudes\n",
    "amp_cmap = ListedColormap([(255/255,253/255,255/255),(238/255,212/255,0/255),(238/255,212/255,0/255),\n",
    "                         (225/255,182/255,0/255), (225/255,182/255,0/255),\n",
    "                         (213/255,151/255,0/255), (213/255,151/255,0/255),\n",
    "                         (201/255,121/255,0/255), (201/255,121/255,0/255),\n",
    "                         (189/255,91/255,0/255),(189/255,91/255,0/255),(177/255,60/255,0/255),(177/255,60/255,0/255),\n",
    "                         (165/255,30/255,0/255),(165/255,30/255,0/255),(153/255,0/255,0/255),(153/255,0/255,0/255),])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be1572f0-68af-4fad-bb5a-6526908eed8e",
   "metadata": {},
   "source": [
    "The orientation model head outputs the sine and cosine of the wave orientation angle. In order to visualise these orientations as vectors, we will use the `quiver` method in matplotlib. This takes vector components as input rather than the angle, so we need to write a function to compute these vector components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c9274b-fc21-489f-9ec7-80f09d6ceaed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute horizontal and vertical components of orientation vectors\n",
    "def quiver_orient(dataset, sep=32, xcoord='projection_x_coordinate', ycoord='projection_y_coordinate'):\n",
    "    angle_rad = dataset['orientation'].values*(np.pi/180) # convert degrees to radians\n",
    "    \n",
    "    new_x = np.zeros(int(512/sep)) # define coarser coordinate grid\n",
    "    new_y = np.zeros(int(512/sep))\n",
    "    angle_rad2 = np.zeros((int(512/sep), int(512/sep))) # initialize coarse matrix of orientation angles\n",
    "\n",
    "    # Populate matrix by iterating over it and computing angles\n",
    "    i = 0\n",
    "    while i < len(angle_rad2):\n",
    "        j = 0\n",
    "        while j < len(angle_rad2[i]):\n",
    "            new_angle_rad = np.pi/2 - angle_rad[sep*i][sep*j] # write orientation angles between 0 and pi\n",
    "            angle_rad2[i][j] = new_angle_rad # store orientation angle in angle_rad2\n",
    "            new_y[j] = dataset[ycoord][j*sep] # store y coordinate\n",
    "            j = j + 1\n",
    "        new_x[i] = dataset[xcoord][i*sep] # store y coordinate\n",
    "        i = i + 1\n",
    "\n",
    "    # Store coarse array in xarray Dataset\n",
    "    alt_dataframe = xr.Dataset(\n",
    "        data_vars={'angle_rad': ((ycoord, xcoord), angle_rad2)},\n",
    "        coords = {xcoord: new_x, ycoord: new_y}\n",
    "    )\n",
    "    sf = 2 # scale factor to make arrows larger on plot\n",
    "\n",
    "    # Compute horizontal and vertical components of scaled orientation vectors\n",
    "    alt_dataframe['orient_u'] = ((ycoord, xcoord), sf*np.cos(angle_rad2))\n",
    "    alt_dataframe['orient_v'] = ((ycoord, xcoord), sf*np.sin(angle_rad2))\n",
    "    alt_dataframe['-orient_u'] = ((ycoord, xcoord), sf*-np.cos(angle_rad2))\n",
    "    alt_dataframe['-orient_v'] = ((ycoord, xcoord), sf*-np.sin(angle_rad2))\n",
    "    return alt_dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a28de6-8be9-48b2-a1e8-ba90153b600f",
   "metadata": {},
   "source": [
    "We are now ready to define our main plotting function. This takes an xarray `DataSet` containing the input and output of our models, and produces a four-panel figure showing the model predictions. Collecting the plotting code into a function like this reduces pollution of our global namespace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8197515e-df54-4d3d-9608-b4bda10fe9b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to plot the model predictions\n",
    "def plot(ds, data='ukv'):\n",
    "    if data == 'ukv':\n",
    "        with open('data/projection/crs.pkl', 'rb') as projfile:\n",
    "            proj = pickle.load(projfile) # import coordinate reference system for Met Office data\n",
    "        xcoord = 'projection_x_coordinate'\n",
    "        ycoord = 'projection_y_coordinate'\n",
    "    if data == 'synthetic':\n",
    "        proj = None # synthetic data on regular x-y grid\n",
    "        xcoord = 'x'\n",
    "        ycoord = 'y'\n",
    "    \n",
    "    fig = plt.figure(figsize=(13, 10), layout='constrained')\n",
    "    ax1 = fig.add_subplot(221, projection=proj)\n",
    "    ax2 = fig.add_subplot(222, projection=proj)\n",
    "    ax3 = fig.add_subplot(223, projection=proj)\n",
    "    ax4 = fig.add_subplot(224, projection=proj)\n",
    "    \n",
    "    ds['upward_air_velocity'].plot.pcolormesh(\n",
    "        cmap=vv_cmap,\n",
    "        robust=False,\n",
    "        rasterized=True,\n",
    "        ax=ax1,\n",
    "        vmin=-4.25,\n",
    "        vmax=4.25,\n",
    "        add_colorbar=True,\n",
    "        cbar_kwargs={\n",
    "            'label':'Upward Air Velocity (m s $^{-1}$)',\n",
    "            'shrink':0.6,\n",
    "            'ticks':np.arange(-4,5,1),\n",
    "            'extend':'neither'\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    ds['segmentation'].plot.contour(\n",
    "        cmap=ListedColormap(['black']),\n",
    "        alpha=1,\n",
    "        add_colorbar=False,\n",
    "        ax=ax1\n",
    "    )\n",
    "    \n",
    "    ds['wavelength'].plot.pcolormesh(\n",
    "        cmap='viridis',\n",
    "        alpha=1,\n",
    "        add_colorbar=True,\n",
    "        rasterized=True,\n",
    "        ax=ax2,\n",
    "        cbar_kwargs={\n",
    "            'label':'Wavelength (km)',\n",
    "            'shrink':0.6,\n",
    "            'extend':'neither'\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    ds['amplitude'].plot.pcolormesh(\n",
    "        cmap=amp_cmap,\n",
    "        vmin=0,\n",
    "        vmax=4.25,\n",
    "        robust=False,\n",
    "        rasterized=True,\n",
    "        ax=ax3,\n",
    "        add_colorbar=True,\n",
    "        cbar_kwargs={\n",
    "            'label':'Amplitude Prediction (m s $^{-1}$)',\n",
    "            'shrink':0.6,'ticks':np.arange(0,4.5,0.5),\n",
    "            'extend':'neither'\n",
    "        }\n",
    "    )\n",
    "            \n",
    "    ds['upward_air_velocity'].plot.pcolormesh(\n",
    "        cmap=vv_cmap,\n",
    "        robust=False,\n",
    "        rasterized=True,\n",
    "        ax=ax4,\n",
    "        vmin=-4.25,\n",
    "        vmax=4.25,\n",
    "        add_colorbar=True,\n",
    "        alpha=1,\n",
    "        cbar_kwargs={\n",
    "            'label':'Upward Air Velocity (m s $^{-1}$)',\n",
    "            'shrink':0.6,\n",
    "            'ticks':np.arange(-4,5,1),\n",
    "            'extend':'neither'\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    headlength = 3\n",
    "    headaxislength = 2 # set these to 0 for no arrows\n",
    "\n",
    "    width = 0.004 # arrow width\n",
    "    ds2 = quiver_orient(ds, sep=16, xcoord=xcoord, ycoord=ycoord) # dataset with horizontal and vertical components\n",
    "    \n",
    "    ds2.plot.quiver(\n",
    "        xcoord, ycoord, 'orient_u', 'orient_v',\n",
    "        ax=ax4,\n",
    "        transform=proj,\n",
    "        width=width,\n",
    "        pivot='tail',\n",
    "        headlength=headlength,\n",
    "        headaxislength=headaxislength,\n",
    "        add_guide=False\n",
    "    )\n",
    "\n",
    "    ds2.plot.quiver(\n",
    "        xcoord, ycoord, '-orient_u', '-orient_v',\n",
    "        ax=ax4,\n",
    "        transform=proj,\n",
    "        width=width,\n",
    "        pivot='tail',\n",
    "        headlength=headlength,\n",
    "        headaxislength=headaxislength,\n",
    "        add_guide=False\n",
    "    ) # add arrows pointing in opposite direction\n",
    "    \n",
    "    ax1.set_title('700 hPa Vertical Velocity and segmentation mask')\n",
    "    ax2.set_title('Wavelength')\n",
    "    ax3.set_title('Amplitude')\n",
    "    ax4.set_title('Orientation (perpendicular to wave fronts)')\n",
    "    \n",
    "    if proj != None:\n",
    "        for ax in [ax1, ax2, ax3, ax4]:\n",
    "            ax.coastlines('10m', alpha=0.5)\n",
    "    if data == 'ukv':\n",
    "        forecast_time = str(ds['forecast_reference_time'].values)[:-10] + 'Z'\n",
    "        fig.suptitle('Lee Wave Data: Characteristics Prediction ' + forecast_time)\n",
    "    if data == 'synthetic':\n",
    "        fig.suptitle('Synthetic Wave Characteristic Prediction', y=.93)\n",
    "            \n",
    "    plt.savefig('model_predictions.pdf', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "102af206-23d2-4f1f-9df0-6366ca1dd8b2",
   "metadata": {},
   "source": [
    "Finally, we are ready to plot our results!\n",
    "\n",
    "We start by loading in test data from over the UK on the 14th of February 2021."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b240009d-0e29-47e0-874b-4b5ea5a53755",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in test data for plotting\n",
    "leewaves = xr.open_dataset('data/test_feb/vertical_velocities/20210214T0900Z-PT0000H00M-wind_vertical_velocity_at_700hPa.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d64220f-4481-47d4-8634-5d5ab3e0a768",
   "metadata": {},
   "source": [
    "Next, we read in our trained models and make our predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e26971cf-31fd-4191-bd1e-aabc5d6d44be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load trained models\n",
    "models = load_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b406c27c-12c8-4128-90d7-e308121e6c51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make predictions with trained models\n",
    "output = predict(models, leewaves)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "555a3019-358c-4ef4-b676-ba5bfe916238",
   "metadata": {},
   "source": [
    "Calling our plotting function shows the results here and also saves to disc as a pdf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f35657d4-5495-4c28-95f3-fd842f38b880",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the model predictions and save output to disc\n",
    "plot(output, data='ukv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dda2ae2-52d1-4ec6-bebe-d9094c7cb306",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}