srgan_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Enhanced Super-Resolution Generative Adversarial Network training**\n",
    "\n",
    "Here in this jupyter notebook, we will train an adapted Enhanced Super-Resolution Generative Adversarial Network (ESRGAN),\n",
    "to create a high-resolution (250m) Antarctic bed Digital Elevation Model(DEM) from a low-resolution (1000m) BEDMAP2 DEM.\n",
    "In addition to that, we use additional correlated inputs that can also tell us something about the bed topography.\n",
    "\n",
    "<img src=\"https://yuml.me/diagram/scruffy;dir:LR/class/[Antarctic Snow Accumulation (1000m)]->[Generator model],[MEASURES Ice Flow Velocity (450m)]->[Generator model],[REMA (100m)]->[Generator model],[BEDMAP2 (1000m)]->[Generator model],[Generator model]->[High res bed DEM (250m)],[High res bed DEM (250m)]->[Discriminator model],[Groundtruth Image (250m)]->[Discriminator model],[Discriminator model]->[True/False]\" alt=\"4 input ESRGAN model\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0. Setup libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python       : 3.7.3 | packaged by conda-forge | (default, Jul  1 2019, 21:52:21) \n",
      "Platform: Linux-4.15.0-46-generic-x86_64-with-debian-buster-sid\n",
      "Chainer: 7.0.0rc1\n",
      "ChainerX: Not Available\n",
      "NumPy: 1.17.0rc2\n",
      "CuPy:\n",
      "  CuPy Version          : 7.0.0rc1\n",
      "  CUDA Root             : /usr/local/cuda\n",
      "  CUDA Build Version    : 10000\n",
      "  CUDA Driver Version   : 10020\n",
      "  CUDA Runtime Version  : 10000\n",
      "  cuDNN Build Version   : 7301\n",
      "  cuDNN Version         : 7301\n",
      "  NCCL Build Version    : 2402\n",
      "  NCCL Runtime Version  : 2402\n",
      "iDeep: Not Available\n"
     ]
    }
   ],
   "source": [
    "import functools\n",
    "import glob\n",
    "import os\n",
    "import random\n",
    "import shutil\n",
    "import sys\n",
    "import typing\n",
    "\n",
    "import comet_ml\n",
    "import IPython.display\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pygmt as gmt\n",
    "import quilt\n",
    "import rasterio\n",
    "import tqdm\n",
    "import xarray as xr\n",
    "\n",
    "import chainer\n",
    "import chainer.functions as F\n",
    "import chainer.links as L\n",
    "import cupy\n",
    "import livelossplot\n",
    "import optuna\n",
    "import ssim.functions\n",
    "\n",
    "from features.environment import _load_ipynb_modules, _download_model_weights_from_comet\n",
    "\n",
    "try:  # check if CUDA_VISIBLE_DEVICES environment variable is set\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"]\n",
    "except KeyError:  # if not set, then set it to the first GPU\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "print(\"Python       :\", sys.version.split(\"\\n\")[0])\n",
    "chainer.print_runtime_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chainer configurations https://docs.chainer.org/en/latest/reference/configuration.html\n",
    "# Use CuDNN deterministic mode\n",
    "chainer.global_config.cudnn_deterministic = True\n",
    "\n",
    "# Set seed values\n",
    "seed = 42\n",
    "random.seed = seed\n",
    "np.random.seed(seed=seed)\n",
    "if cupy.is_available():\n",
    "    for c in range(cupy.cuda.runtime.getDeviceCount()):\n",
    "        with cupy.cuda.Device(c):\n",
    "            cupy.random.seed(seed=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load data\n",
    "- Download pre-packaged data from [Quilt](https://github.com/quiltdata/quilt)\n",
    "- Convert arrays for Chainer, from Numpy (CPU) to CuPy (GPU) format (if available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data_into_memory(\n",
    "    refresh_cache: bool = True,\n",
    "    quilt_hash: str = \"580960bc97696f7ca89dba61fb6225a2ff631d49876fefef8dc05a033f13e14f\",\n",
    ") -> (chainer.datasets.dict_dataset.DictDataset, str):\n",
    "    \"\"\"\n",
    "    Downloads the prepackaged tiled data from quilt based on a hash,\n",
    "    and loads it into CPU or GPU memory depending on what is available.\n",
    "    \"\"\"\n",
    "\n",
    "    if refresh_cache:\n",
    "        quilt.install(\n",
    "            package=\"weiji14/deepbedmap/model/train\", hash=quilt_hash, force=True\n",
    "        )\n",
    "    pkg = quilt.load(pkginfo=\"weiji14/deepbedmap/model/train\", hash=quilt_hash)\n",
    "\n",
    "    W1_data = pkg.W1_data()  # miscellaneous data REMA\n",
    "    W2_data = pkg.W2_data()  # miscellaneous data MEASURES Ice Flow\n",
    "    W3_data = pkg.W3_data()  # miscellaneous data Arthern Accumulation\n",
    "    X_data = pkg.X_data()  # low resolution BEDMAP2\n",
    "    Y_data = pkg.Y_data()  # high resolution groundtruth\n",
    "    # print(W1_data.shape, W2_data.shape, W3_data.shape, X_data.shape, Y_data.shape)\n",
    "\n",
    "    # Detect if there is a CUDA GPU first\n",
    "    if cupy.is_available():\n",
    "        print(\"Using GPU\")\n",
    "        W1_data = chainer.backend.cuda.to_gpu(array=W1_data, device=None)\n",
    "        W2_data = chainer.backend.cuda.to_gpu(array=W2_data, device=None)\n",
    "        W3_data = chainer.backend.cuda.to_gpu(array=W3_data, device=None)\n",
    "        X_data = chainer.backend.cuda.to_gpu(array=X_data, device=None)\n",
    "        Y_data = chainer.backend.cuda.to_gpu(array=Y_data, device=None)\n",
    "    else:\n",
    "        print(\"Using CPU only\")\n",
    "\n",
    "    return (\n",
    "        chainer.datasets.DictDataset(\n",
    "            X=X_data, W1=W1_data, W2=W2_data, W3=W3_data, Y=Y_data\n",
    "        ),\n",
    "        quilt_hash,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Split dataset into training (train) and development (dev) sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_dev_iterators(\n",
    "    dataset: chainer.datasets.dict_dataset.DictDataset,\n",
    "    first_size: int,  # size of training set\n",
    "    batch_size: int = 128,\n",
    "    seed: int = 42,\n",
    ") -> (\n",
    "    chainer.iterators.serial_iterator.SerialIterator,\n",
    "    int,\n",
    "    chainer.iterators.serial_iterator.SerialIterator,\n",
    "    int,\n",
    "):\n",
    "    \"\"\"\n",
    "    Create Chainer Dataset Iterators after splitting dataset into\n",
    "    training and development (validation) sets.\n",
    "    \"\"\"\n",
    "\n",
    "    # Train/Dev split of the dataset\n",
    "    train_set, dev_set = chainer.datasets.split_dataset_random(\n",
    "        dataset=dataset, first_size=first_size, seed=seed\n",
    "    )\n",
    "\n",
    "    # Create Chainer Dataset Iterators out of the split datasets\n",
    "    train_iter = chainer.iterators.SerialIterator(\n",
    "        dataset=train_set, batch_size=batch_size, repeat=True, shuffle=True\n",
    "    )\n",
    "    dev_iter = chainer.iterators.SerialIterator(\n",
    "        dataset=dev_set, batch_size=batch_size, repeat=True, shuffle=False\n",
    "    )\n",
    "\n",
    "    print(\n",
    "        f\"Training dataset: {len(train_set)} tiles,\",\n",
    "        f\"Development dataset: {len(dev_set)} tiles\",\n",
    "    )\n",
    "\n",
    "    return train_iter, len(train_set), dev_iter, len(dev_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Architect model\n",
    "\n",
    "Enhanced Super Resolution Generative Adversarial Network (ESRGAN) model based on [Wang et al. 2018](https://arxiv.org/abs/1809.00219v2).\n",
    "Refer to original Pytorch implementation at https://github.com/xinntao/ESRGAN.\n",
    "See also previous (non-enhanced) SRGAN model architecture by [Ledig et al. 2017](https://arxiv.org/abs/1609.04802)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Generator Network Architecture\n",
    "\n",
    "![ESRGAN architecture - Generator Network composed of many Dense Convolutional Blocks](https://github.com/xinntao/ESRGAN/raw/master/figures/architecture.jpg)\n",
    "\n",
    "3 main components: 1) Input Block, 2) Residual Blocks, 3) Upsampling Blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.1 Input block, specially customized for DeepBedMap to take in 3 different inputs\n",
    "\n",
    "Details of the first convolutional layer for each input:\n",
    "\n",
    "- Input tiles are 11000m by 11000m.\n",
    "- Convolution filter kernels are 3000m by 3000m.\n",
    "- Strides are 1000m by 1000m.\n",
    "\n",
    "Example: for a 100m spatial resolution tile:\n",
    "\n",
    "- Input tile is 110pixels by 110pixels\n",
    "- Convolution filter kernels are 30pixels by 30pixels\n",
    "- Strides are 10pixels by 10pixels\n",
    "\n",
    "Note that the first convolutional layers uses **valid** padding, see https://github.com/weiji14/deepbedmap/pull/65."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepbedmapInputBlock(chainer.Chain):\n",
    "    \"\"\"\n",
    "    Custom input block for DeepBedMap.\n",
    "    Takes in BEDMAP2 (X), REMA Ice Surface Elevation (W1),\n",
    "    MEaSUREs Ice Surface Velocity x and y components (W2) and Snow Accumulation (W3).\n",
    "    Passes them through custom-sized convolutions, with the results being concatenated.\n",
    "\n",
    "    Each filter kernel is 3km by 3km in size, with a 1km stride and no padding.\n",
    "    So for a 1km resolution image, (i.e. 1km pixel size):\n",
    "    kernel size is (3, 3), stride is (1, 1), and pad is (0, 0)\n",
    "\n",
    "    X  (?,1,11,11)  --Conv2D-- (?,32,9,9) \\\n",
    "    W1 (?,1,110,110) --Conv2D-- (?,32,9,9) --Concat-- (?,128,9,9)\n",
    "    W2 (?,2,22,22)  --Conv2D-- (?,32,9,9) /\n",
    "    W3 (?,1,11,11) --Conv2D-- (?,32,9,9) /\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, out_channels=32):\n",
    "        super().__init__()\n",
    "        init_weights = chainer.initializers.HeNormal(scale=0.1, fan_option=\"fan_in\")\n",
    "\n",
    "        with self.init_scope():\n",
    "            self.conv_on_X = L.Convolution2D(\n",
    "                in_channels=1,\n",
    "                out_channels=out_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=(0, 0),  # 'valid' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_on_W1 = L.Convolution2D(\n",
    "                in_channels=1,\n",
    "                out_channels=out_channels,\n",
    "                ksize=(30, 30),\n",
    "                stride=(10, 10),\n",
    "                pad=(0, 0),  # 'valid' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_on_W2 = L.Convolution2D(\n",
    "                in_channels=2,\n",
    "                out_channels=out_channels,\n",
    "                ksize=(6, 6),\n",
    "                stride=(2, 2),\n",
    "                pad=(0, 0),  # 'valid' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_on_W3 = L.Convolution2D(\n",
    "                in_channels=1,\n",
    "                out_channels=out_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=(0, 0),  # 'valid' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "\n",
    "    def forward(self, x, w1, w2, w3):\n",
    "        \"\"\"\n",
    "        Forward computation, i.e. evaluate based on inputs X, W1, W2 and W3\n",
    "        \"\"\"\n",
    "        x_ = self.conv_on_X(x)\n",
    "        w1_ = self.conv_on_W1(w1)\n",
    "        w2_ = self.conv_on_W2(w2)\n",
    "        w3_ = self.conv_on_W3(w3)\n",
    "\n",
    "        a = F.concat(xs=(x_, w1_, w2_, w3_))\n",
    "        return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.2 Residual Block\n",
    "\n",
    "![The Residual in Residual Dense Block in detail](https://raw.githubusercontent.com/xinntao/ESRGAN/master/figures/RRDB.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualDenseBlock(chainer.Chain):\n",
    "    \"\"\"\n",
    "    Residual Dense Block made up of 5 Convolutional2D-LeakyReLU layers.\n",
    "    Final output has a residual scaling factor.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_out_channels: int = 64,\n",
    "        inter_channels: int = 32,\n",
    "        residual_scaling: float = 0.1,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.residual_scaling = residual_scaling\n",
    "        init_weights = chainer.initializers.HeNormal(scale=0.1, fan_option=\"fan_in\")\n",
    "\n",
    "        with self.init_scope():\n",
    "            self.conv_layer1 = L.Convolution2D(\n",
    "                in_channels=in_out_channels,\n",
    "                out_channels=inter_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_layer2 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=inter_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_layer3 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=inter_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_layer4 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=inter_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_layer5 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=in_out_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward computation, i.e. evaluate based on input x\n",
    "        \"\"\"\n",
    "        a0 = x\n",
    "\n",
    "        a1 = self.conv_layer1(a0)\n",
    "        a1 = F.leaky_relu(x=a1, slope=0.2)\n",
    "        a1_cat = F.concat(xs=(a0, a1), axis=1)\n",
    "\n",
    "        a2 = self.conv_layer2(a1_cat)\n",
    "        a2 = F.leaky_relu(x=a2, slope=0.2)\n",
    "        a2_cat = F.concat(xs=(a0, a1, a2), axis=1)\n",
    "\n",
    "        a3 = self.conv_layer3(a2_cat)\n",
    "        a3 = F.leaky_relu(x=a3, slope=0.2)\n",
    "        a3_cat = F.concat(xs=(a0, a1, a2, a3), axis=1)\n",
    "\n",
    "        a4 = self.conv_layer4(a3_cat)\n",
    "        a4 = F.leaky_relu(x=a4, slope=0.2)\n",
    "        a4_cat = F.concat(xs=(a0, a1, a2, a3, a4), axis=1)\n",
    "\n",
    "        a5 = self.conv_layer5(a4_cat)\n",
    "\n",
    "        # Final concatenation, with residual scaling of 0.2\n",
    "        a6 = F.add(a5 * self.residual_scaling, a0)\n",
    "\n",
    "        return a6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResInResDenseBlock(chainer.Chain):\n",
    "    \"\"\"\n",
    "    Residual in Residual Dense block made of 3 Residual Dense Blocks\n",
    "\n",
    "       ------------  ----------  ------------\n",
    "      |            ||          ||            |\n",
    "    -----DenseBlock--DenseBlock--DenseBlock-(+)--\n",
    "      |                                      |\n",
    "       --------------------------------------\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, denseblock_class=ResidualDenseBlock, residual_scaling: float = 0.1\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.residual_scaling = residual_scaling\n",
    "\n",
    "        with self.init_scope():\n",
    "            self.residual_dense_block1 = denseblock_class(\n",
    "                residual_scaling=residual_scaling\n",
    "            )\n",
    "            self.residual_dense_block2 = denseblock_class(\n",
    "                residual_scaling=residual_scaling\n",
    "            )\n",
    "            self.residual_dense_block3 = denseblock_class(\n",
    "                residual_scaling=residual_scaling\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward computation, i.e. evaluate based on input x\n",
    "        \"\"\"\n",
    "        a1 = self.residual_dense_block1(x)\n",
    "        a2 = self.residual_dense_block2(a1)\n",
    "        a3 = self.residual_dense_block3(a2)\n",
    "\n",
    "        # Final concatenation, with residual scaling of 0.2\n",
    "        a4 = F.add(a3 * self.residual_scaling, x)\n",
    "\n",
    "        return a4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.3 Build the Generator Network, with upsampling layers!\n",
    "\n",
    "![4 inputs feeding into the Generator Network, producing a high resolution prediction output](https://yuml.me/dfd301a2.png)\n",
    "\n",
    "<!--\n",
    "[W3_input(ACCUMULATION)|1x11x11]-k3n32s1>[W3_inter|32x9x9],[W3_inter]->[Concat|128x9x9]\n",
    "[W2_input(MEASURES)|2x22x22]-k6n32s2>[W2_inter|32x9x9],[W2_inter]->[Concat|128x9x9]\n",
    "[W1_input(REMA)|1x110x110]-k30n32s10>[W1_inter|32x9x9],[W1_inter]->[Concat|128x9x9]\n",
    "[X_input(BEDMAP2)|1x11x11]-k3n32s1>[X_inter|32x9x9],[X_inter]->[Concat|128x9x9]\n",
    "[Concat|128x9x9]->[Generator-Network|Many-Residual-Blocks],[Generator-Network]->[Y_hat(High-Resolution_DEM)|1x36x36]\n",
    "-->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GeneratorModel(chainer.Chain):\n",
    "    \"\"\"\n",
    "    The generator network which is a deconvolutional neural network.\n",
    "    Converts a low resolution input into a super resolution output.\n",
    "\n",
    "    Glues the input block with several residual blocks and upsampling layers\n",
    "\n",
    "    Parameters:\n",
    "      num_residual_blocks -- how many Residual-in-Residual Dense Blocks to use\n",
    "      residual_scaling -- scale factor for residuals before adding to parent branch\n",
    "      out_channels -- integer representing number of output channels/filters/kernels\n",
    "\n",
    "    Example:\n",
    "      A convolved input_shape of (9,9,1) passing through b residual blocks with\n",
    "      a scaling of 4 and out_channels 1 will result in an image of shape (36,36,1)\n",
    "\n",
    "    >>> generator_model = GeneratorModel()\n",
    "    >>> y_pred = generator_model.forward(\n",
    "    ...     x=np.random.rand(1, 1, 11, 11).astype(\"float32\"),\n",
    "    ...     w1=np.random.rand(1, 1, 110, 110).astype(\"float32\"),\n",
    "    ...     w2=np.random.rand(1, 2, 22, 22).astype(\"float32\"),\n",
    "    ...     w3=np.random.rand(1, 1, 11, 11).astype(\"float32\"),\n",
    "    ... )\n",
    "    >>> y_pred.shape\n",
    "    (1, 1, 36, 36)\n",
    "    >>> generator_model.count_params()\n",
    "    8907749\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        inblock_class=DeepbedmapInputBlock,\n",
    "        resblock_class=ResInResDenseBlock,\n",
    "        num_residual_blocks: int = 12,\n",
    "        residual_scaling: float = 0.1,\n",
    "        out_channels: int = 1,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.num_residual_blocks = num_residual_blocks\n",
    "        self.residual_scaling = residual_scaling\n",
    "        init_weights = chainer.initializers.HeNormal(scale=0.1, fan_option=\"fan_in\")\n",
    "\n",
    "        with self.init_scope():\n",
    "\n",
    "            # Initial Input and Residual Blocks\n",
    "            self.input_block = inblock_class()\n",
    "            self.pre_residual_conv_layer = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.residual_network = resblock_class(\n",
    "                residual_scaling=residual_scaling\n",
    "            ).repeat(n_repeat=num_residual_blocks)\n",
    "            self.post_residual_conv_layer = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "\n",
    "            # Upsampling Layers\n",
    "            self.post_upsample_conv_layer_1 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.post_upsample_conv_layer_2 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "\n",
    "            # Final post-upsample convolution layers\n",
    "            self.final_conv_layer1 = L.DeformableConvolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                offset_initialW=init_weights,\n",
    "                deform_initialW=init_weights,\n",
    "            )\n",
    "            self.final_conv_layer2 = L.DeformableConvolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=out_channels,\n",
    "                ksize=(3, 3),\n",
    "                stride=(1, 1),\n",
    "                pad=1,  # 'same' padding\n",
    "                offset_initialW=init_weights,\n",
    "                deform_initialW=init_weights,\n",
    "            )\n",
    "\n",
    "    def forward(\n",
    "        self, x: cupy.ndarray, w1: cupy.ndarray, w2: cupy.ndarray, w3: cupy.ndarray\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Forward computation, i.e. evaluate based on input tensors\n",
    "\n",
    "        Each input should be either a numpy or cupy array.\n",
    "        \"\"\"\n",
    "        # 0 part\n",
    "        # Resize inputs to right scale using convolution\n",
    "        # with hardcoded kernel_sizes, strides and padding lengths\n",
    "        # Also concatenate all inputs\n",
    "        a0 = self.input_block(x=x, w1=w1, w2=w2, w3=w3)\n",
    "\n",
    "        # 1st part\n",
    "        # Pre-residual k3n64s1\n",
    "        a1 = self.pre_residual_conv_layer(a0)\n",
    "        a1 = F.leaky_relu(x=a1, slope=0.2)\n",
    "\n",
    "        # 2nd part\n",
    "        # Residual blocks k3n64s1\n",
    "        a2 = self.residual_network(a1)\n",
    "\n",
    "        # 3rd part\n",
    "        # Post-residual blocks k3n64s1\n",
    "        a3 = self.post_residual_conv_layer(a2)\n",
    "        a3 = F.add(a1, a3)\n",
    "\n",
    "        # 4th part\n",
    "        # Upsampling (hardcoded to be 4x, actually 2x run twice)\n",
    "        # Uses Nearest Neighbour Interpolation followed by Convolution2D k3n64s1\n",
    "        a4_1 = F.resize_images(\n",
    "            x=a3, output_shape=(2 * a3.shape[-2], 2 * a3.shape[-1]), mode=\"nearest\"\n",
    "        )\n",
    "        a4_1 = self.post_upsample_conv_layer_1(a4_1)\n",
    "        a4_1 = F.leaky_relu(x=a4_1, slope=0.2)\n",
    "\n",
    "        a4_2 = F.resize_images(\n",
    "            x=a4_1,\n",
    "            output_shape=(2 * a4_1.shape[-2], 2 * a4_1.shape[-1]),\n",
    "            mode=\"nearest\",\n",
    "        )\n",
    "        a4_2 = self.post_upsample_conv_layer_2(a4_2)\n",
    "        a4_2 = F.leaky_relu(x=a4_2, slope=0.2)\n",
    "\n",
    "        # 5th part\n",
    "        # Generate high resolution output k3n64s1 and k3n1s1\n",
    "        a5_1 = self.final_conv_layer1(a4_2)\n",
    "        a5_1 = F.leaky_relu(x=a5_1, slope=0.2)\n",
    "        a5_2 = self.final_conv_layer2(a5_1)\n",
    "\n",
    "        return a5_2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Discriminator Network Architecture\n",
    "\n",
    "Discriminator implementation following that of [ESRGAN](https://arxiv.org/abs/1809.00219).\n",
    "[VGG-style](https://arxiv.org/abs/1409.1556)\n",
    "Consists of 10 Conv2D-BatchNorm-LeakyReLU blocks, followed by 2 Fully Connected Layers of size 100 and 1, with **no** final sigmoid activation.\n",
    "Note also how the BatchNormalization layers **are still preserved**.\n",
    "Original Pytorch implementation can be found [here](https://github.com/xinntao/BasicSR/blame/902b4ae1f4beec7359de6e62ed0aebfc335d8dfd/codes/models/modules/architecture.py#L86-L129).\n",
    "\n",
    "![Discriminator Network](https://yuml.me/diagram/scruffy/class/[High-Resolution_DEM|32x32x1]->[Discriminator-Network],[Discriminator-Network]->[False/True|0/1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiscriminatorModel(chainer.Chain):\n",
    "    \"\"\"\n",
    "    The discriminator network which is a convolutional neural network.\n",
    "    Takes ONE high resolution input image and predicts whether it is\n",
    "    real or fake on a scale of 0 to 1, where 0 is fake and 1 is real.\n",
    "\n",
    "    Consists of several Conv2D-BatchNorm-LeakyReLU blocks, followed by\n",
    "    a fully connected linear layer with LeakyReLU activation and a final\n",
    "    fully connected linear layer with Sigmoid activation.\n",
    "\n",
    "    >>> discriminator_model = DiscriminatorModel()\n",
    "    >>> y_pred = discriminator_model.forward(\n",
    "    ...     x=np.random.rand(2, 1, 36, 36).astype(\"float32\")\n",
    "    ... )\n",
    "    >>> y_pred.shape\n",
    "    (2, 1)\n",
    "    >>> discriminator_model.count_params()\n",
    "    10370761\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        init_weights = chainer.initializers.HeNormal(scale=0.1, fan_option=\"fan_in\")\n",
    "\n",
    "        with self.init_scope():\n",
    "\n",
    "            self.conv_layer0 = L.Convolution2D(\n",
    "                in_channels=None,\n",
    "                out_channels=64,\n",
    "                ksize=3,\n",
    "                stride=1,\n",
    "                pad=1,  # 'same' padding\n",
    "                nobias=False,  # only first Conv2D layer uses bias\n",
    "                initialW=init_weights,\n",
    "            )\n",
    "            self.conv_layer1 = L.Convolution2D(None, 64, 4, 2, 1, True, init_weights)\n",
    "            self.conv_layer2 = L.Convolution2D(None, 128, 3, 1, 1, True, init_weights)\n",
    "            self.conv_layer3 = L.Convolution2D(None, 128, 4, 2, 1, True, init_weights)\n",
    "            self.conv_layer4 = L.Convolution2D(None, 128, 3, 1, 1, True, init_weights)\n",
    "            self.conv_layer5 = L.Convolution2D(None, 256, 4, 2, 1, True, init_weights)\n",
    "            self.conv_layer6 = L.Convolution2D(None, 256, 3, 1, 1, True, init_weights)\n",
    "            self.conv_layer7 = L.Convolution2D(None, 512, 4, 2, 1, True, init_weights)\n",
    "            self.conv_layer8 = L.Convolution2D(None, 512, 3, 1, 1, True, init_weights)\n",
    "            self.conv_layer9 = L.Convolution2D(None, 512, 4, 2, 1, True, init_weights)\n",
    "\n",
    "            self.batch_norm1 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm2 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm3 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm4 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm5 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm6 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm7 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm8 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "            self.batch_norm9 = L.BatchNormalization(axis=(0, 2, 3), eps=1e-5)\n",
    "\n",
    "            self.linear_1 = L.Linear(in_size=None, out_size=100, initialW=init_weights)\n",
    "            self.linear_2 = L.Linear(in_size=None, out_size=1, initialW=init_weights)\n",
    "\n",
    "    def forward(self, x: cupy.ndarray):\n",
    "        \"\"\"\n",
    "        Forward computation, i.e. evaluate based on input tensor\n",
    "\n",
    "        Each input should be either a numpy or cupy array.\n",
    "        \"\"\"\n",
    "\n",
    "        # 1st part\n",
    "        # Convolutonal Block without Batch Normalization k3n64s1\n",
    "        a0 = self.conv_layer0(x=x)\n",
    "        a0 = F.leaky_relu(x=a0, slope=0.2)\n",
    "\n",
    "        # 2nd part\n",
    "        # Convolutional Blocks with Batch Normalization k3n{64*f}s{1or2}\n",
    "        a1 = self.conv_layer1(x=a0)\n",
    "        a1 = self.batch_norm1(x=a1)\n",
    "        a1 = F.leaky_relu(x=a1, slope=0.2)\n",
    "        a2 = self.conv_layer2(x=a1)\n",
    "        a2 = self.batch_norm2(x=a2)\n",
    "        a2 = F.leaky_relu(x=a2, slope=0.2)\n",
    "        a3 = self.conv_layer3(x=a2)\n",
    "        a3 = self.batch_norm3(x=a3)\n",
    "        a3 = F.leaky_relu(x=a3, slope=0.2)\n",
    "        a4 = self.conv_layer4(x=a3)\n",
    "        a4 = self.batch_norm4(x=a4)\n",
    "        a4 = F.leaky_relu(x=a4, slope=0.2)\n",
    "        a5 = self.conv_layer5(x=a4)\n",
    "        a5 = self.batch_norm5(x=a5)\n",
    "        a5 = F.leaky_relu(x=a5, slope=0.2)\n",
    "        a6 = self.conv_layer6(x=a5)\n",
    "        a6 = self.batch_norm6(x=a6)\n",
    "        a6 = F.leaky_relu(x=a6, slope=0.2)\n",
    "        a7 = self.conv_layer7(x=a6)\n",
    "        a7 = self.batch_norm7(x=a7)\n",
    "        a7 = F.leaky_relu(x=a7, slope=0.2)\n",
    "        a8 = self.conv_layer8(x=a7)\n",
    "        a8 = self.batch_norm8(x=a8)\n",
    "        a8 = F.leaky_relu(x=a8, slope=0.2)\n",
    "        a9 = self.conv_layer9(x=a8)\n",
    "        a9 = self.batch_norm9(x=a9)\n",
    "        a9 = F.leaky_relu(x=a9, slope=0.2)\n",
    "\n",
    "        # 3rd part\n",
    "        # Flatten, Dense (Fully Connected) Layers and Output\n",
    "        a10 = F.reshape(x=a9, shape=(len(a9), -1))  # flatten while keeping batch_size\n",
    "        a10 = self.linear_1(x=a10)\n",
    "        a10 = F.leaky_relu(x=a10, slope=0.2)\n",
    "        a11 = self.linear_2(x=a10)\n",
    "        # a11 = F.sigmoid(x=a11)  # no sigmoid activation, as it is in the loss function\n",
    "\n",
    "        return a11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Define Loss function and Metrics for the Generator and Discriminator Networks\n",
    "\n",
    "Now we define the Perceptual Loss function for our Generator and Discriminator neural network models, where:\n",
    "\n",
    "$$Perceptual Loss = Content Loss + Adversarial Loss + Topographic Loss + Structural Loss$$\n",
    "\n",
    "![Perceptual Loss in an adapted Enhanced Super Resolution Generative Adversarial Network](https://yuml.me/19155033.png)\n",
    "\n",
    "<!--\n",
    "[LowRes-Inputs]-Generator>[SuperResolution_DEM]\n",
    "[SuperResolution_DEM]-.->[note:Content-Loss|MeanAbsoluteError{bg:yellow}]\n",
    "[LowRes-Inputs]-.->[note:Topographic-Loss|MeanAbsoluteError{bg:yellow}]\n",
    "[SuperResolution_DEM]-.->[note:Structural-Loss|SSIM{bg:yellow}]\n",
    "[SuperResolution_DEM]-.->[note:Topographic-Loss]\n",
    "[HighRes-Groundtruth_DEM]-.->[note:Content-Loss]\n",
    "[HighRes-Groundtruth_DEM]-.->[note:Structural-Loss]\n",
    "[SuperResolution_DEM]-Discriminator>[False_or_True_Prediction]\n",
    "[HighRes-Groundtruth_DEM]-Discriminator>[False_or_True_Prediction]\n",
    "[False_or_True_Prediction]<->[False_or_True_Label]\n",
    "[False_or_True_Prediction]-.->[note:Adversarial-Loss|BinaryCrossEntropy{bg:yellow}]\n",
    "[False_or_True_Label]-.->[note:Adversarial-Loss]\n",
    "[note:Content-Loss]-.->[note:Perceptual-Loss{bg:gold}]\n",
    "[note:Adversarial-Loss]-.->[note:Perceptual-Loss{bg:gold}]\n",
    "[note:Topographic-Loss]-.->[note:Perceptual-Loss{bg:gold}]\n",
    "[note:Structural-Loss]-.->[note:Perceptual-Loss{bg:gold}]\n",
    "-->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Content Loss\n",
    "\n",
    "The original SRGAN paper by [Ledig et al. 2017](https://arxiv.org/abs/1609.04802v5) calculates *Content Loss* based on the ReLU activation layers of the pre-trained 19 layer VGG network.\n",
    "The implementation below is less advanced, simply using an L1 loss, i.e., a pixel-wise [Mean Absolute Error (MAE) loss](https://docs.chainer.org/en/latest/reference/generated/chainer.functions.mean_absolute_error.html) as the *Content Loss*.\n",
    "Specifically, the *Content Loss* is calculated as the MAE difference between the output of the generator model (i.e. the predicted Super Resolution Image) and that of the groundtruth image (i.e. the true High Resolution Image).\n",
    "\n",
    "$$ e_i = ||G(x_{i}) - y_i||_{1} $$\n",
    "\n",
    "$$ Loss_{Content} = Mean Absolute Error = \\dfrac{1}{n} \\sum\\limits_{i=1}^n e_i $$\n",
    "\n",
    "where $G(x_{i})$ is the Generator Network's predicted value, and $y_i$ is the groundtruth value, respectively at pixel $i$.\n",
    "$e_i$ thus represents the absolute error (L1 loss) (denoted by $||\\dots||_{1}$) between the predicted and groundtruth value.\n",
    "We then sum all the pixel-wise errors $e_i,\\dots,e_n$ and divide by the number of pixels $n$ to get the Arithmetic Mean $\\dfrac{1}{n} \\sum\\limits_{i=1}^n$ of our error which is our *Content Loss*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adversarial Loss\n",
    "\n",
    "The *Adversarial Loss* or *Generative Loss* (confusing I know) is the same as in the original SRGAN paper.\n",
    "It is defined based on the probabilities of the discriminator believing that the reconstructed Super Resolution Image is a natural High Resolution Image.\n",
    "The implementation below uses the [Binary CrossEntropy loss](https://keras.io/losses/#binary_crossentropy).\n",
    "Specifically, this *Adversarial Loss* is calculated between the output of the discriminator model (a value between 0 and 1) and that of the groundtruth label (a boolean value of either 0 or 1).\n",
    "\n",
    "$$ Loss_{Adversarial} = Binary Cross Entropy Loss = -\\dfrac{1}{n} \\sum\\limits_{i=1}^n ( y_i ln(\\sigma(x_i)) + (1-y_i) ln(1 - \\sigma(x_i) ) $$\n",
    "\n",
    "where $\\sigma$ is the [Sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function) activation function, $\\sigma = \\dfrac{1}{1+e^{-x}} = \\dfrac{e^x}{e^x+1}$, $y_i$ is the groundtruth label (1 for real, 0 for fake) and $x_i$ is the prediction (before sigmoid activation is applied), all respectively at pixel $i$.\n",
    "\n",
    "$\\sigma(x)$ is basically the sigmoid activated output from a Standard Discriminator neural network, which some people also denote as $D(.)$.\n",
    "Technically, some people also write $D(x) = \\sigma(C(x))$, where $C(x)$ is the raw, non-transformed output from the Discriminator neural network (i.e. no sigmoid activation applied) on the input data $x$.\n",
    "For simplicity, we now denote $C(x)$ simply as $x$ in the following equations, i.e. using $\\sigma(x)$ to replace $\\sigma(C(x))$.\n",
    "\n",
    "Again, the [Binary Cross Entropy Loss](https://en.wikipedia.org/wiki/Cross_entropy#Cross-entropy_error_function_and_logistic_regression) calculated on one pixel is defined as follows:\n",
    "\n",
    "$$ -( y ln(\\sigma(x)) + (1-y) ln(1 - \\sigma(x) )$$\n",
    "\n",
    "With the full expansion as such:\n",
    "\n",
    "$$ -\\bigg[ y ln\\big(\\dfrac{e^x}{e^x+1}\\big) + (1-y) ln\\big(1 - \\dfrac{e^x}{e^x+1}\\big) \\bigg] $$\n",
    "\n",
    "The above equation is mathematically equivalent to the one below, and can be derived using [Logarithm rules](https://en.wikipedia.org/wiki/Logarithm#Product,_quotient,_power,_and_root) such as the Power Rule and Product Rule, and using the fact that $ln(e)=1$ and $ln(1)=0$:\n",
    "\n",
    "$$ -[ xy - ln(1+e^x) ] $$\n",
    "\n",
    "However, this reformed equation is numerically unstable (see discussion [here](https://www.reddit.com/r/MachineLearning/comments/4euzmk/logsumexp_for_logistic_regression/)), and is good for values of $x<0$.\n",
    "For values of $x>=0$, there is an alternative representation which we can derive:\n",
    "\n",
    "$$ -[ xy - ln(1+e^x) - x + x ] $$\n",
    "$$ -[ x(y-1) - ln(1 + e^x) + ln(e^x) ] $$\n",
    "$$ -\\bigg[ x(y-1) - ln\\big(\\dfrac{e^x}{1+e^x}\\big) \\bigg] $$\n",
    "$$ -\\bigg[ x(y-1) - ln\\big(\\dfrac{1}{1+e^{-x}}\\big) \\bigg] $$\n",
    "$$ - [ x(y-1) - ln(1) + ln(1+e^{-x}) ] $$\n",
    "$$ - [ x(y-1) + ln(1+e^{-x}) $$\n",
    "\n",
    "In order to have a numerically stable function that works for both $x<0$ and $x>=0$, we can write it like so as in Caffe's implementation:\n",
    "\n",
    "$$ -[ x(y - 1_{x>=0} - ln(1+e^{x-2x\\cdot1_{x>=0}}) ] $$\n",
    "\n",
    "Alternatively, Chainer does it like so:\n",
    "\n",
    "$$ -[ x(y - 1_{x>=0} - ln(1+e^{-|x|}) ] $$\n",
    "\n",
    "Or in Python code (the Chainer implemention from [here](https://github.com/chainer/chainer/blob/v6.0.0b1/chainer/functions/loss/sigmoid_cross_entropy.py#L41-L44)), bearing in mind that the natural logarithm $ln$ is `np.log` in Numpy:\n",
    "\n",
    "```python\n",
    "    sigmoidbinarycrossentropyloss = -(x * (y - (x >= 0)) - np.log1p(np.exp(-np.abs(x))))\n",
    "```\n",
    "\n",
    "See also how [Pytorch](https://pytorch.org/docs/stable/nn.html?highlight=bcewithlogitsloss#torch.nn.BCEWithLogitsLoss) and [Tensorflow](https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits) implements this in a numerically stable manner."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Topographic Loss\n",
    "\n",
    "In addition to the L1 Content Loss, we further define a *Topographic Loss*.\n",
    "Specifically, we want each of the averaged value in each 4x4 grid of the predicted DeepBedMap image to correspond to a 1x1 pixel on the BEDMAP2 image.\n",
    "\n",
    "Due to BEDMAP2 having a 4x lower resolution than the predicted DeepBedMap DEM (1000m compared to 250m), we first apply a 4x4 [Mean/Average Pooling](https://docs.chainer.org/en/latest/reference/generated/chainer.functions.average_pooling_2d.html) operation on the DeepBedMap image, turning it into a 1x1 pixel grid that matches the shape of the BEDMAP2 image.\n",
    "\n",
    "$$ \\bar{\\hat{y}}_j = Mean = \\dfrac{1}{n} \\sum\\limits_{i=1}^n y_i $$\n",
    "\n",
    "where $\\bar{\\hat{y}}_j$ is the mean/average of all predicted pixel values $y_i$ across the 16 high resolution (DeepBedMap) pixels $i$ within a 4x4 grid corresponding to the spatial location of one low resolution (BEDMAP2) pixel at position $j$.\n",
    "Then, we calculate the [MAE](https://docs.chainer.org/en/latest/reference/generated/chainer.functions.mean_absolute_error.html) difference between the output of the generator model (i.e. the predicted Super-Resolution DeepBedMap bed DEM Image) and that of the BEDMAP2 (i.e. the original Low Resolution Image we are super-resolving).\n",
    "\n",
    "$$ e_j = ||\\bar{\\hat{y}}_j - x_j||_{1} $$\n",
    "\n",
    "where $\\bar{\\hat{y}}_j$ is the mean of the 4x4 pixels we just calculated above, and $x_j$ is the spatially corresponding BEDMAP2 pixel, respectively at BEDMAP2 pixel $j$.\n",
    "$e_j$ thus represents the absolute error (L1 loss) (denoted by $||\\dots||_{1}$) between the (averaged) super-resolution and low-resolution values.\n",
    "We then sum all the pixel-wise errors $e_j,\\dots,e_m$ and divide by the number of low resolution pixels $m$ to get the Arithmetic Mean $\\dfrac{1}{m} \\sum\\limits_{i=1}^m$ of our error which is our *Topographic Loss*.\n",
    "\n",
    "$$ Loss_{Topographic} = Mean Absolute Error = \\dfrac{1}{m} \\sum\\limits_{j=1}^m e_j $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Structural Loss\n",
    "\n",
    "Complementing the L1 Content Loss further, we define a *Structural (Similarity) Loss*.\n",
    "This is computed using the [Structural Similarity (SSIM)](https://en.wikipedia.org/wiki/Structural_similarity) Index,\n",
    "on a moving window (here set to 9x9) between the super resolution predicted image and high resolution groundtruth image.\n",
    "The comparison takes into account luminance, contrast and structural information and is calculated over a single window patch like so:\n",
    "\n",
    "$$SSIM(x,y) = \\dfrac{(2\\mu_x\\mu_y + c_1)(2\\sigma_{xy} + c_2)}{(\\mu_x^2 + \\mu_y^2 + c_1)(\\sigma_x^2 + \\sigma_y^2 + c_2)} $$\n",
    "\n",
    "where $\\mu_x$ and $\\mu_y$ are the average (mean) of predicted image $x$ and groundtruth image $y$ respectively,\n",
    "$\\sigma_{xy}$ is the covariance of $x$ and $y$,\n",
    "$\\sigma_x^2$ and $\\sigma_y^2$ are the variance of $x$ and $y$ respectively, and\n",
    "$c_1$ and $c_2$ are two variables set to $0.01^2$ and $0.03^2$ to stabilize division with a weak denominator.\n",
    "Our Structural Loss can then be formulated as follows:\n",
    "\n",
    "$$ Loss_{Structural} = 1 - \\dfrac{1}{p} \\sum\\limits_{i=1}^p SSIM(x, y)_p $$\n",
    "\n",
    "where we do $1$ minus the mean of all structural similarity values $SSIM(x, y)$ calculated over every patch $p$ obtained via a sliding window over our predicted image $x$ and groundtruth image $y$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_generator_loss(\n",
    "    y_pred: chainer.variable.Variable,\n",
    "    y_true: cupy.ndarray,\n",
    "    fake_labels: cupy.ndarray,\n",
    "    real_labels: cupy.ndarray,\n",
    "    fake_minus_real_target: cupy.ndarray,\n",
    "    real_minus_fake_target: cupy.ndarray,\n",
    "    x_topo: cupy.ndarray,\n",
    "    content_loss_weighting: float = 1e-2,  # e.g. ~35 * 1e-2 = 0.35\n",
    "    adversarial_loss_weighting: float = 2e-2,  # e.g. ~10 * 2e-2 = 0.20\n",
    "    topographic_loss_weighting: float = 2e-3,  # e.g. ~35 * 2e-3 = 0.07\n",
    "    structural_loss_weighting: float = 5.25e-0,  # e.g. ~0.75 * 5.25e-0 = 3.9375\n",
    ") -> chainer.variable.Variable:\n",
    "    \"\"\"\n",
    "    This function calculates the weighted sum between\n",
    "    \"Content Loss\", \"Adversarial Loss\", \"Topographic Loss\", and \"Structural Loss\"\n",
    "    which forms the basis for training the Generator Network.\n",
    "\n",
    "    >>> calculate_generator_loss(\n",
    "    ...     y_pred=chainer.variable.Variable(data=np.ones(shape=(2, 1, 12, 12))),\n",
    "    ...     y_true=np.full(shape=(2, 1, 12, 12), fill_value=10.0),\n",
    "    ...     fake_labels=np.array([[-1.2], [0.5]]),\n",
    "    ...     real_labels=np.array([[0.5], [-0.8]]),\n",
    "    ...     fake_minus_real_target=np.array([[1], [1]]).astype(np.int32),\n",
    "    ...     real_minus_fake_target=np.array([[0], [0]]).astype(np.int32),\n",
    "    ...     x_topo=np.full(shape=(2, 1, 3, 3), fill_value=9.0),\n",
    "    ... )\n",
    "    variable(4.35108415)\n",
    "    \"\"\"\n",
    "    # Content Loss (L1, Mean Absolute Error) between predicted and groundtruth 2D images\n",
    "    content_loss = F.mean_absolute_error(x0=y_pred, x1=y_true)\n",
    "\n",
    "    # Adversarial Loss between 1D labels\n",
    "    adversarial_loss = calculate_discriminator_loss(\n",
    "        real_labels_pred=real_labels,\n",
    "        fake_labels_pred=fake_labels,\n",
    "        real_minus_fake_target=real_minus_fake_target,  # Zeros (0) instead of ones (1)\n",
    "        fake_minus_real_target=fake_minus_real_target,  # Ones (1) instead of zeros (0)\n",
    "    )\n",
    "\n",
    "    # Topographic Loss (L1, Mean Absolute Error) between predicted and low res 2D images\n",
    "    topographic_loss = F.mean_absolute_error(\n",
    "        x0=F.average_pooling_2d(x=y_pred, ksize=(4, 4)), x1=x_topo\n",
    "    )\n",
    "\n",
    "    # Structural Similarity Loss between predicted and groundtruth 2D images\n",
    "    structural_loss = 1 - ssim_loss_func(y_pred=y_pred, y_true=y_true)\n",
    "\n",
    "    # Get generator loss\n",
    "    weighted_content_loss = content_loss_weighting * content_loss\n",
    "    weighted_adversarial_loss = adversarial_loss_weighting * adversarial_loss\n",
    "    weighted_topographic_loss = topographic_loss_weighting * topographic_loss\n",
    "    weighted_structural_loss = structural_loss_weighting * structural_loss\n",
    "\n",
    "    g_loss = (\n",
    "        weighted_content_loss\n",
    "        + weighted_adversarial_loss\n",
    "        + weighted_topographic_loss\n",
    "        + weighted_structural_loss\n",
    "    )\n",
    "\n",
    "    return g_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def psnr(\n",
    "    y_pred: cupy.ndarray, y_true: cupy.ndarray, data_range=2 ** 32\n",
    ") -> cupy.ndarray:\n",
    "    \"\"\"\n",
    "    Peak Signal-Noise Ratio (PSNR) metric, calculated batchwise.\n",
    "    See https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio#Definition\n",
    "\n",
    "    Can take in either numpy (CPU) or cupy (GPU) arrays as input.\n",
    "    Implementation is same as skimage.measure.compare_psnr with data_range=2**32\n",
    "\n",
    "    >>> psnr(\n",
    "    ...     y_pred=np.ones(shape=(2, 1, 3, 3)),\n",
    "    ...     y_true=np.full(shape=(2, 1, 3, 3), fill_value=2),\n",
    "    ... )\n",
    "    192.65919722494797\n",
    "    \"\"\"\n",
    "    xp = chainer.backend.get_array_module(y_true)\n",
    "\n",
    "    # Calculate Mean Squred Error along predetermined axes\n",
    "    mse = xp.mean(xp.square(xp.subtract(y_pred, y_true)), axis=None)\n",
    "\n",
    "    # Calculate Peak Signal-Noise Ratio, setting MAX_I as 2^32, i.e. max for int32\n",
    "    return xp.multiply(20, xp.log10(data_range / xp.sqrt(mse)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ssim_loss_func(\n",
    "    y_pred: chainer.variable.Variable,\n",
    "    y_true: cupy.ndarray,\n",
    "    window_size: int = 9,\n",
    "    stride: int = 1,\n",
    ") -> chainer.variable.Variable:\n",
    "    \"\"\"\n",
    "    Structural Similarity (SSIM) loss/metric, calculated with default window size of 9.\n",
    "    See https://en.wikipedia.org/wiki/Structural_similarity\n",
    "\n",
    "    Can take in either numpy (CPU) or cupy (GPU) arrays as input.\n",
    "\n",
    "    >>> ssim_loss_func(\n",
    "    ...     y_pred=chainer.variable.Variable(data=np.ones(shape=(2, 1, 9, 9))),\n",
    "    ...     y_true=np.full(shape=(2, 1, 9, 9), fill_value=2.0),\n",
    "    ... )\n",
    "    variable(0.800004)\n",
    "    \"\"\"\n",
    "    if not y_pred.shape == y_true.shape:\n",
    "        raise ValueError(\"Input images must have the same dimensions.\")\n",
    "\n",
    "    ssim_value = ssim.functions.ssim_loss(\n",
    "        y=y_pred, t=y_true, window_size=window_size, stride=stride\n",
    "    )\n",
    "    return ssim_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_discriminator_loss(\n",
    "    real_labels_pred: chainer.variable.Variable,\n",
    "    fake_labels_pred: chainer.variable.Variable,\n",
    "    real_minus_fake_target: cupy.ndarray,\n",
    "    fake_minus_real_target: cupy.ndarray,\n",
    ") -> chainer.variable.Variable:\n",
    "    \"\"\"\n",
    "    This function purely calculates the \"Adversarial Loss\"\n",
    "    in a Relativistic Average Generative Adversarial Network (RaGAN).\n",
    "\n",
    "    It forms the basis for training the Discriminator Network,\n",
    "    but it is also used as part of the Generator Network's loss function.\n",
    "\n",
    "    See paper by Jolicoeur-Martineau, 2018 at https://arxiv.org/abs/1807.00734\n",
    "    for the mathematical details of the RaGAN loss function.\n",
    "\n",
    "    Original Sigmoid_Cross_Entropy formula:\n",
    "    -(y * np.log(sigmoid(x)) + (1 - y) * np.log(1 - sigmoid(x)))\n",
    "\n",
    "    Numerically stable formula:\n",
    "    -(x * (y - (x >= 0)) - np.log1p(np.exp(-np.abs(x))))\n",
    "\n",
    "    where y = the target difference between real and fake labels (i.e. 1 - 0 = 1)\n",
    "          x = the calculated difference between real_labels_pred and fake_labels_pred\n",
    "\n",
    "    >>> calculate_discriminator_loss(\n",
    "    ...     real_labels_pred=chainer.variable.Variable(data=np.array([[1.1], [-0.5]])),\n",
    "    ...     fake_labels_pred=chainer.variable.Variable(data=np.array([[-0.3], [1.0]])),\n",
    "    ...     real_minus_fake_target=np.array([[1], [1]]),\n",
    "    ...     fake_minus_real_target=np.array([[0], [0]]),\n",
    "    ... )\n",
    "    variable(1.56670504)\n",
    "    \"\"\"\n",
    "\n",
    "    # Calculate arithmetic mean of real/fake predicted labels\n",
    "    real_labels_pred_avg = F.mean(real_labels_pred)\n",
    "    fake_labels_pred_avg = F.mean(fake_labels_pred)\n",
    "\n",
    "    # Binary Cross-Entropy Loss with Sigmoid\n",
    "    real_versus_fake_loss = F.sigmoid_cross_entropy(\n",
    "        x=(real_labels_pred - fake_labels_pred_avg), t=real_minus_fake_target\n",
    "    )  # let predicted labels from real images be more realistic than those from fake\n",
    "    fake_versus_real_loss = F.sigmoid_cross_entropy(\n",
    "        x=(fake_labels_pred - real_labels_pred_avg), t=fake_minus_real_target\n",
    "    )  # let predicted labels from fake images be less realistic than those from real\n",
    "\n",
    "    # Relativistic average Standard GAN's Discriminator Loss\n",
    "    d_loss = real_versus_fake_loss + fake_versus_real_loss\n",
    "\n",
    "    return d_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the models\n",
    "def compile_srgan_model(\n",
    "    num_residual_blocks: int = 12,\n",
    "    residual_scaling: float = 0.1,\n",
    "    learning_rate: float = 1.6e-4,\n",
    "):\n",
    "    \"\"\"\n",
    "    Instantiate our Super Resolution Generative Adversarial Network (SRGAN) model here.\n",
    "    The Generator and Discriminator neural networks are created,\n",
    "    and an Adam loss optimization function is linked to the models.\n",
    "\n",
    "    Returns:\n",
    "    1) generator_model\n",
    "    2) generator_optimizer\n",
    "    3) discriminator_model\n",
    "    4) discriminator_optimizer\n",
    "    \"\"\"\n",
    "\n",
    "    # Instantiate our Generator and Discriminator Neural Network models\n",
    "    generator_model = GeneratorModel(\n",
    "        num_residual_blocks=num_residual_blocks, residual_scaling=residual_scaling\n",
    "    )\n",
    "    discriminator_model = DiscriminatorModel()\n",
    "\n",
    "    # Transfer models to GPU if available\n",
    "    if cupy.is_available():  # Check if CuPy was loaded, i.e. GPU is available\n",
    "        generator_model.to_gpu(device=None)\n",
    "        discriminator_model.to_gpu(device=None)\n",
    "\n",
    "    # Setup optimizer, using Adam\n",
    "    generator_optimizer = chainer.optimizers.Adam(alpha=learning_rate, eps=1e-8).setup(\n",
    "        link=generator_model\n",
    "    )\n",
    "    discriminator_optimizer = chainer.optimizers.Adam(\n",
    "        alpha=learning_rate, eps=1e-8\n",
    "    ).setup(link=discriminator_model)\n",
    "\n",
    "    return (\n",
    "        generator_model,\n",
    "        generator_optimizer,\n",
    "        discriminator_model,\n",
    "        discriminator_optimizer,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Train model\n",
    "\n",
    "[Gherkin](https://en.wikipedia.org/wiki/Gherkin_(language))/Plain English statement at what the Super-Resolution Generative Adversarial Network below does\n",
    "\n",
    "```gherkin\n",
    "    # language: en\n",
    "    Feature: SRGAN DeepBedMap\n",
    "      In order to create a great map of Antarctica's bed\n",
    "      As a data scientist,\n",
    "      We want a model that produces realistic images from many open datasets\n",
    "\n",
    "      Scenario: Train discriminator to beat generator\n",
    "        Given fake generated images from a generator\n",
    "          And real groundtruth images\n",
    "         When the two sets of images are fed into the discriminator for comparison\n",
    "         Then the discriminator should learn to know the fakes from the real images\n",
    "\n",
    "      Scenario: Train generator to fool discriminator\n",
    "        Given fake generated images from a generator\n",
    "          And what we think the discriminator believes is real\n",
    "         When we compare the fake images to the real ones\n",
    "         Then the generator should learn to create a more authentic looking image\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_eval_discriminator(\n",
    "    input_arrays: typing.Dict[str, cupy.ndarray],\n",
    "    g_model,\n",
    "    d_model,\n",
    "    d_optimizer=None,\n",
    "    train: bool = True,\n",
    ") -> (float, float):\n",
    "    \"\"\"\n",
    "    Trains the Discriminator within a Super Resolution Generative Adversarial Network.\n",
    "    Discriminator is trainable, Generator is not trained (only produces predictions).\n",
    "\n",
    "    Steps:\n",
    "    - Generator produces fake images\n",
    "    - Fake images combined with real groundtruth images\n",
    "    - Discriminator trained with these images and their Fake(0)/Real(1) labels\n",
    "\n",
    "    >>> train_arrays = {\n",
    "    ...     \"X\": np.random.RandomState(seed=42).rand(2, 1, 11, 11).astype(np.float32),\n",
    "    ...     \"W1\": np.random.RandomState(seed=42).rand(2, 1, 110, 110).astype(np.float32),\n",
    "    ...     \"W2\": np.random.RandomState(seed=42).rand(2, 2, 22, 22).astype(np.float32),\n",
    "    ...     \"W3\": np.random.RandomState(seed=42).rand(2, 1, 11, 11).astype(np.float32),\n",
    "    ...     \"Y\": np.random.RandomState(seed=42).rand(2, 1, 36, 36).astype(np.float32),\n",
    "    ... }\n",
    "    >>> discriminator_model = DiscriminatorModel()\n",
    "    >>> discriminator_optimizer = chainer.optimizers.Adam(alpha=0.001, eps=1e-7).setup(\n",
    "    ...     link=discriminator_model\n",
    "    ... )\n",
    "    >>> generator_model = GeneratorModel()\n",
    "\n",
    "    >>> d_weight0 = [d for d in discriminator_model.params()][-3][0].array\n",
    "    >>> d_train_loss, d_train_accu = train_eval_discriminator(\n",
    "    ...     input_arrays=train_arrays,\n",
    "    ...     g_model=generator_model,\n",
    "    ...     d_model=discriminator_model,\n",
    "    ...     d_optimizer=discriminator_optimizer,\n",
    "    ... )\n",
    "    >>> d_weight1 = [d for d in discriminator_model.params()][-3][0].array\n",
    "    >>> d_weight0 != d_weight1  #check that training has occurred (i.e. weights changed)\n",
    "    True\n",
    "    \"\"\"\n",
    "    # @pytest.fixture\n",
    "    chainer.global_config.train = train  # Explicitly set Chainer's train/eval flag\n",
    "    if train is True:\n",
    "        assert d_optimizer is not None  # Optimizer required for neural network training\n",
    "    xp = chainer.backend.get_array_module(input_arrays[\"Y\"])\n",
    "\n",
    "    # @given(\"fake generated images from a generator\")\n",
    "    with chainer.using_config(name=\"enable_backprop\", value=False):\n",
    "        fake_images = g_model.forward(\n",
    "            x=input_arrays[\"X\"],\n",
    "            w1=input_arrays[\"W1\"],\n",
    "            w2=input_arrays[\"W2\"],\n",
    "            w3=input_arrays[\"W3\"],\n",
    "        ).array\n",
    "    fake_labels = xp.zeros(shape=(len(fake_images), 1)).astype(xp.int32)\n",
    "\n",
    "    # @given(\"real groundtruth images\")\n",
    "    real_images = input_arrays[\"Y\"]\n",
    "    real_labels = xp.ones(shape=(len(real_images), 1)).astype(xp.int32)\n",
    "\n",
    "    # @when(\"the two sets of images are fed into the discriminator for comparison\")\n",
    "    real_labels_pred = d_model.forward(x=real_images)\n",
    "    fake_labels_pred = d_model.forward(x=fake_images)\n",
    "    real_minus_fake_target = xp.ones(shape=(len(real_images), 1)).astype(xp.int32)\n",
    "    fake_minus_real_target = xp.zeros(shape=(len(real_images), 1)).astype(xp.int32)\n",
    "    d_loss = calculate_discriminator_loss(\n",
    "        real_labels_pred=real_labels_pred,  # real image should get close to 1\n",
    "        fake_labels_pred=fake_labels_pred,  # fake image should get close to 0\n",
    "        real_minus_fake_target=real_minus_fake_target,  # where 1 (real) - 0 (fake) = 1 (target)\n",
    "        fake_minus_real_target=fake_minus_real_target,  # where 0 (fake) - 1 (real) = 0 (target)?\n",
    "    )\n",
    "\n",
    "    predicted_labels = xp.concatenate([real_labels_pred.array, fake_labels_pred.array])\n",
    "    groundtruth_labels = xp.concatenate([real_labels, fake_labels])\n",
    "    d_accu = F.binary_accuracy(y=predicted_labels, t=groundtruth_labels)\n",
    "\n",
    "    # @then(\"the discriminator should learn to know the fakes from the real images\")\n",
    "    if train is True:\n",
    "        d_model.cleargrads()  # clear/zero all gradients\n",
    "        d_loss.backward()  # renew gradients\n",
    "        d_optimizer.update()  # backpropagate the loss using optimizer\n",
    "\n",
    "    return float(d_loss.array), float(d_accu.array)  # return discriminator metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_eval_generator(\n",
    "    input_arrays: typing.Dict[str, cupy.ndarray],\n",
    "    g_model,\n",
    "    d_model,\n",
    "    g_optimizer=None,\n",
    "    train: bool = True,\n",
    ") -> (float, float, float):\n",
    "    \"\"\"\n",
    "    Evaluates and/or trains the Generator for one minibatch\n",
    "    within a Super Resolution Generative Adversarial Network.\n",
    "    Discriminator is not trainable, Generator is trained.\n",
    "\n",
    "    If train is set to False, only forward pass is run, i.e. evaluation/prediction only\n",
    "    If train is set to True, forward and backward pass are run, i.e. train with backprop\n",
    "\n",
    "    Steps:\n",
    "    - Generator produces fake images\n",
    "    - Fake images compared with real groundtruth images\n",
    "    - Generator is trained to be more like real image\n",
    "\n",
    "    >>> train_arrays = {\n",
    "    ...     \"X\": np.random.RandomState(seed=42).rand(2, 1, 11, 11).astype(np.float32),\n",
    "    ...     \"W1\": np.random.RandomState(seed=42).rand(2, 1, 110, 110).astype(np.float32),\n",
    "    ...     \"W2\": np.random.RandomState(seed=42).rand(2, 2, 22, 22).astype(np.float32),\n",
    "    ...     \"W3\": np.random.RandomState(seed=42).rand(2, 1, 11, 11).astype(np.float32),\n",
    "    ...     \"Y\": np.random.RandomState(seed=42).rand(2, 1, 36, 36).astype(np.float32),\n",
    "    ... }\n",
    "    >>> generator_model = GeneratorModel()\n",
    "    >>> generator_optimizer = chainer.optimizers.Adam(alpha=0.001, eps=1e-7).setup(\n",
    "    ...     link=generator_model\n",
    "    ... )\n",
    "    >>> discriminator_model = DiscriminatorModel()\n",
    "\n",
    "    >>> g_weight0 = [g for g in generator_model.params()][8][0, 0, 0, 0].array\n",
    "    >>> _ = train_eval_generator(\n",
    "    ...     input_arrays=train_arrays,\n",
    "    ...     g_model=generator_model,\n",
    "    ...     d_model=discriminator_model,\n",
    "    ...     g_optimizer=generator_optimizer,\n",
    "    ... )\n",
    "    >>> g_weight1 = [g for g in generator_model.params()][8][0, 0, 0, 0].array\n",
    "    >>> g_weight0 != g_weight1  #check that training has occurred (i.e. weights changed)\n",
    "    True\n",
    "    \"\"\"\n",
    "\n",
    "    # @pytest.fixture\n",
    "    chainer.global_config.train = train  # Explicitly set Chainer's train/eval flag\n",
    "    if train is True:\n",
    "        assert g_optimizer is not None  # Optimizer required for neural network training\n",
    "    xp = chainer.backend.get_array_module(input_arrays[\"Y\"])\n",
    "\n",
    "    # @given(\"fake generated images from a generator\")\n",
    "    fake_images = g_model.forward(\n",
    "        x=input_arrays[\"X\"],\n",
    "        w1=input_arrays[\"W1\"],\n",
    "        w2=input_arrays[\"W2\"],\n",
    "        w3=input_arrays[\"W3\"],\n",
    "    )\n",
    "    with chainer.using_config(name=\"train\", value=False):  # Using non-train BatchNorm\n",
    "        fake_labels = d_model.forward(x=fake_images).array.astype(xp.float32)\n",
    "\n",
    "    # @given(\"what we think the discriminator believes is real\")\n",
    "    real_images = input_arrays[\"Y\"]\n",
    "    real_labels = xp.ones(shape=(len(real_images), 1)).astype(xp.float32)\n",
    "\n",
    "    # @when(\"we compare the fake images to the real ones\")\n",
    "    fake_minus_real_target = xp.ones(shape=(len(real_images), 1)).astype(xp.int32)\n",
    "    real_minus_fake_target = xp.zeros(shape=(len(real_images), 1)).astype(xp.int32)\n",
    "    g_loss = calculate_generator_loss(\n",
    "        # content loss inputs, 2D images\n",
    "        y_pred=fake_images,\n",
    "        y_true=real_images,\n",
    "        # adversarial loss inputs, 1D labels\n",
    "        fake_labels=fake_labels,  # fake label 'should' get close to 1\n",
    "        real_labels=real_labels,  # real label 'should' get close to 0\n",
    "        fake_minus_real_target=fake_minus_real_target,  # where 1 (fake) - 0 (real) = 1 (target)\n",
    "        real_minus_fake_target=real_minus_fake_target,  # where 0 (real) - 1 (fake) = 0 (target)?\n",
    "        # topographic loss inputs, 2D image of low resolution input\n",
    "        x_topo=input_arrays[\"X\"][:, :, 1:-1, 1:-1],  # sliced to remove 1km borders\n",
    "    )\n",
    "    g_psnr = psnr(y_pred=fake_images.array, y_true=real_images)\n",
    "    g_ssim = ssim_loss_func(y_pred=fake_images, y_true=real_images)\n",
    "\n",
    "    # @then(\"the generator should learn to create a more authentic looking image\")\n",
    "    if train is True:\n",
    "        g_model.cleargrads()  # clear/zero all gradients\n",
    "        g_loss.backward()  # renew gradients\n",
    "        g_optimizer.update()  # backpropagate the loss using optimizer\n",
    "\n",
    "    return (\n",
    "        float(g_loss.array),\n",
    "        float(g_psnr),\n",
    "        float(g_ssim.array),\n",
    "    )  # return generator loss and metric values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainer(\n",
    "    i: int,  # current epoch\n",
    "    columns: list,  # dataframe column names, i.e. the metric names\n",
    "    train_iter: chainer.iterators.serial_iterator.SerialIterator,\n",
    "    dev_iter: chainer.iterators.serial_iterator.SerialIterator,\n",
    "    g_model,  # generator_model\n",
    "    g_optimizer,  # generator_optimizer\n",
    "    d_model,  # discriminator_model\n",
    "    d_optimizer,  # discriminator_optimizer\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Trains the Super Resolution Generative Adversarial Networks (SRGAN)'s\n",
    "    Discriminator and Generator components one after another for one epoch.\n",
    "    Also does evaluation on a development dataset and reports metrics.\n",
    "    \"\"\"\n",
    "\n",
    "    metrics_dict = {mn: [] for mn in columns}  # reset metrics dictionary\n",
    "\n",
    "    ## Part 1 - Training on training dataset\n",
    "    while i == train_iter.epoch:  # while we are in epoch i, run minibatch training\n",
    "        train_batch = train_iter.next()\n",
    "        train_arrays = chainer.dataset.concat_examples(batch=train_batch)\n",
    "        ## 1.1 - Train Discriminator\n",
    "        d_train_loss, d_train_accu = train_eval_discriminator(\n",
    "            input_arrays=train_arrays,\n",
    "            g_model=g_model,\n",
    "            d_model=d_model,\n",
    "            d_optimizer=d_optimizer,\n",
    "        )\n",
    "        metrics_dict[\"discriminator_loss\"].append(d_train_loss)\n",
    "        metrics_dict[\"discriminator_accu\"].append(d_train_accu)\n",
    "\n",
    "        ## 1.2 - Train Generator\n",
    "        g_train_loss, g_train_psnr, g_train_ssim = train_eval_generator(\n",
    "            input_arrays=train_arrays,\n",
    "            g_model=g_model,\n",
    "            d_model=d_model,\n",
    "            g_optimizer=g_optimizer,\n",
    "        )\n",
    "        metrics_dict[\"generator_loss\"].append(g_train_loss)\n",
    "        metrics_dict[\"generator_psnr\"].append(g_train_psnr)\n",
    "        metrics_dict[\"generator_ssim\"].append(g_train_ssim)\n",
    "\n",
    "    ## Part 2 - Evaluation on development dataset\n",
    "    while i == dev_iter.epoch:  # while we are in epoch i, evaluate on each minibatch\n",
    "        dev_batch = dev_iter.next()\n",
    "        dev_arrays = chainer.dataset.concat_examples(batch=dev_batch)\n",
    "        ## 2.1 - Evaluate Discriminator\n",
    "        d_dev_loss, d_dev_accu = train_eval_discriminator(\n",
    "            input_arrays=dev_arrays, g_model=g_model, d_model=d_model, train=False\n",
    "        )\n",
    "        metrics_dict[\"val_discriminator_loss\"].append(d_dev_loss)\n",
    "        metrics_dict[\"val_discriminator_accu\"].append(d_dev_accu)\n",
    "\n",
    "        ## 2.2 - Evaluate Generator\n",
    "        g_dev_loss, g_dev_psnr, g_dev_ssim = train_eval_generator(\n",
    "            input_arrays=dev_arrays, g_model=g_model, d_model=d_model, train=False\n",
    "        )\n",
    "        metrics_dict[\"val_generator_loss\"].append(g_dev_loss)\n",
    "        metrics_dict[\"val_generator_psnr\"].append(g_dev_psnr)\n",
    "        metrics_dict[\"val_generator_ssim\"].append(g_dev_ssim)\n",
    "\n",
    "    return metrics_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model_weights_and_architecture(\n",
    "    generator_model, discriminator_model, save_path: str = \"model/weights\"\n",
    ") -> (str, str, str):\n",
    "    \"\"\"\n",
    "    Save the trained neural network's parameter weights and architecture (computational\n",
    "    graph) respectively to zipped Numpy (.npz) and Graphviz DOT (.dot) format.\n",
    "\n",
    "    >>> g_model = GeneratorModel(num_residual_blocks=1)\n",
    "    >>> d_model = DiscriminatorModel()\n",
    "    >>> _, _, _ = save_model_weights_and_architecture(\n",
    "    ...     generator_model=g_model, discriminator_model=d_model, save_path=\"/tmp/weights\"\n",
    "    ... )\n",
    "    >>> os.path.exists(path=\"/tmp/weights/srgan_generator_model_architecture.dot\")\n",
    "    True\n",
    "    \"\"\"\n",
    "\n",
    "    os.makedirs(name=save_path, exist_ok=True)\n",
    "\n",
    "    # Save generator/discriminator model's parameter weights in Numpy Zipped format\n",
    "    generator_model_weights_path: str = os.path.join(\n",
    "        save_path, f\"srgan_generator_model_weights.npz\"\n",
    "    )\n",
    "    chainer.serializers.save_npz(file=generator_model_weights_path, obj=generator_model)\n",
    "    discriminator_model_weights_path: str = os.path.join(\n",
    "        save_path, f\"srgan_discriminator_model_weights.npz\"\n",
    "    )\n",
    "    chainer.serializers.save_npz(\n",
    "        file=discriminator_model_weights_path, obj=discriminator_model\n",
    "    )\n",
    "\n",
    "    # Save generator model's architecture in Graphviz DOT format\n",
    "    model_architecture_path: str = os.path.join(\n",
    "        save_path, f\"srgan_generator_model_architecture.dot\"\n",
    "    )\n",
    "    args = {\n",
    "        \"x\": generator_model.xp.random.rand(128, 1, 11, 11).astype(\"float32\"),\n",
    "        \"w1\": generator_model.xp.random.rand(128, 1, 110, 110).astype(\"float32\"),\n",
    "        \"w2\": generator_model.xp.random.rand(128, 2, 22, 22).astype(\"float32\"),\n",
    "        \"w3\": generator_model.xp.random.rand(128, 1, 11, 11).astype(\"float32\"),\n",
    "    }\n",
    "    graph = chainer.computational_graph.build_computational_graph(\n",
    "        outputs=generator_model.forward(**args)\n",
    "    )\n",
    "    with open(file=model_architecture_path, mode=\"w\") as outgraph:\n",
    "        outgraph.writelines([f\"{line};\\n\" for line in graph.dump().split(\";\")])\n",
    "\n",
    "    return (\n",
    "        generator_model_weights_path,\n",
    "        discriminator_model_weights_path,\n",
    "        model_architecture_path,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Evaluate model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation on independent test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "@functools.lru_cache()\n",
    "def get_fixed_test_inputs(\n",
    "    test_filepath: str = \"highres/20xx_Antarctica_DC8_THW\",\n",
    "    indexers: dict = {},  # {\"y\": slice(1, -2), \"x\": slice(1, -2)},  # custom index-based crop\n",
    ") -> (xr.DataArray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, pd.DataFrame):\n",
    "    \"\"\"\n",
    "    Cached way of getting the fixed (non-model) data array inputs for our hardcoded\n",
    "    DeepBedMap test area along the main trunk of Pine Island Glacier. Use it by simply\n",
    "    calling get_fixed_test_inputs(), which works arounds functools.lru_cache not\n",
    "    accepting unhashable inputs like Python list, dict, and slice objects.\n",
    "    \"\"\"\n",
    "    deepbedmap = _load_ipynb_modules(\"deepbedmap.ipynb\")\n",
    "\n",
    "    # Get neural network input datasets associated with groundtruth image bounds\n",
    "    groundtruth = deepbedmap.get_image_with_bounds(\n",
    "        filepaths=[f\"{test_filepath}.nc\"], indexers=indexers\n",
    "    )\n",
    "    X_tile, W1_tile, W2_tile, W3_tile = deepbedmap.get_deepbedmap_model_inputs(\n",
    "        window_bound=rasterio.coords.BoundingBox(*groundtruth.bounds)\n",
    "    )\n",
    "\n",
    "    # Load xyz points table for test region\n",
    "    data_prep = _load_ipynb_modules(\"data_prep.ipynb\")\n",
    "    points = data_prep.ascii_to_xyz(pipeline_file=f\"{test_filepath}.json\")\n",
    "\n",
    "    return groundtruth, X_tile, W1_tile, W2_tile, W3_tile, points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_deepbedmap_test_result(\n",
    "    model=None,\n",
    "    test_filepath: str = \"highres/20xx_Antarctica_DC8_THW\",\n",
    "    indexers: dict = {},  # {\"y\": slice(1, -2), \"x\": slice(1, -2)},  # custom index-based crop\n",
    "    model_weights_path: str = \"model/weights/srgan_generator_model_weights.npz\",\n",
    ") -> (float, xr.DataArray):\n",
    "    \"\"\"\n",
    "    Gets Root Mean Squared Error of elevation difference between DeepBedMap topography\n",
    "    and reference groundtruth xyz tracks at a particular test region.\n",
    "\n",
    "    If a trained model is passed in, the 'model_weights_path' parameter will be ignored\n",
    "    (i.e. the trained weights will not be reloaded). If instead there is no model passed\n",
    "    in, the default model will be loaded with trained weights from 'model_weights_path'.\n",
    "    \"\"\"\n",
    "    # Get neural network input datasets associated with groundtruth image bounds\n",
    "    # Load xyz points table for test region\n",
    "    groundtruth, X_tile, W1_tile, W2_tile, W3_tile, points = get_fixed_test_inputs()\n",
    "\n",
    "    # Run input datasets through trained neural network model\n",
    "    if model is None:\n",
    "        deepbedmap = _load_ipynb_modules(\"deepbedmap.ipynb\")\n",
    "        model, _ = deepbedmap.load_trained_model(model_weights_path=model_weights_path)\n",
    "    with chainer.using_config(name=\"enable_backprop\", value=False):\n",
    "        Y_hat = model.forward(\n",
    "            x=model.xp.asarray(a=X_tile),\n",
    "            w1=model.xp.asarray(a=W1_tile),\n",
    "            w2=model.xp.asarray(a=W2_tile),\n",
    "            w3=model.xp.asarray(a=W3_tile),\n",
    "        ).array\n",
    "\n",
    "    # Create xarray grid from model prediction\n",
    "    grid = xr.DataArray(\n",
    "        data=np.flipud(cupy.asnumpy(Y_hat[0, 0, :, :])),\n",
    "        dims=[\"y\", \"x\"],\n",
    "        coords={\"y\": groundtruth.y, \"x\": groundtruth.x},\n",
    "    )\n",
    "\n",
    "    # Get the elevation (z) value at specified x, y points along the groundtruth track\n",
    "    df_deepbedmap3 = gmt.grdtrack(points=points, grid=grid, newcolname=\"z_interpolated\")\n",
    "\n",
    "    # Calculate elevation error between groundtruth xyz tracks and deepbedmap\n",
    "    df_deepbedmap3[\"error\"] = df_deepbedmap3.z_interpolated - df_deepbedmap3.z\n",
    "    rmse_deepbedmap3 = (df_deepbedmap3.error ** 2).mean() ** 0.5\n",
    "\n",
    "    return float(rmse_deepbedmap3), grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Hyperparameter tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tuning the various hyperparameters on our test area using [Optuna](https://github.com/pfnet/optuna).\n",
    "Yes, not exactly proper I know, but we have lots of areas we can test on later.\n",
    "\n",
    "Also logging all the experiments using [Comet.ML](https://www.comet.ml) to https://www.comet.ml/weiji14/deepbedmap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def objective(\n",
    "    trial: optuna.trial.Trial = optuna.trial.FixedTrial(\n",
    "        params={\n",
    "            \"batch_size_exponent\": 7,\n",
    "            \"num_residual_blocks\": 12,\n",
    "            \"residual_scaling\": 0.1,\n",
    "            \"learning_rate\": 1.6e-4,\n",
    "            \"num_epochs\": 120,\n",
    "        }\n",
    "    ),\n",
    "    enable_livelossplot: bool = False,  # Default: False, no plots makes it go faster!\n",
    "    enable_comet_logging: bool = True,  # Default: True, log experiment to Comet.ML\n",
    "    resume_experiment_key: str = \"055b697548e048b78202cfebb78d6d8c\",  # Default: None\n",
    "    reload_d_model_weights: bool = False,  # Default: False, Reload discriminator model weights\n",
    ") -> float:\n",
    "    \"\"\"\n",
    "    Objective function for tuning the Hyperparameters of our DeepBedMap model.\n",
    "    Uses the Optuna (https://github.com/pfnet/optuna) library.\n",
    "\n",
    "    List of hyperparameters tuned:\n",
    "    - Learning rate\n",
    "    - Number of residual blocks\n",
    "    - Batch Size\n",
    "    - Number of training epochs\n",
    "    - Residual Scaling factor\n",
    "    \"\"\"\n",
    "\n",
    "    # Start tracking experiment using Comet.ML\n",
    "    experiment = comet_ml.Experiment(\n",
    "        workspace=\"weiji14\",\n",
    "        project_name=\"deepbedmap\",\n",
    "        disabled=not enable_comet_logging,\n",
    "    )\n",
    "    base_model_weight_path: str = f\"model/weights/{experiment.get_key()}\"\n",
    "\n",
    "    # Don't use cached stuff if it's a FixedTrial or the first trial\n",
    "    if not hasattr(trial, \"number\") or trial.number == 0:\n",
    "        refresh_cache = True\n",
    "    elif trial.number > 0:  # Use cache if this is not the first trial\n",
    "        refresh_cache = False\n",
    "\n",
    "    ## Load Dataset\n",
    "    dataset, quilt_hash = load_data_into_memory(refresh_cache=refresh_cache)\n",
    "    experiment.log_parameter(name=\"dataset_hash\", value=quilt_hash)\n",
    "    experiment.log_parameter(name=\"use_gpu\", value=cupy.is_available())\n",
    "    batch_size: int = int(\n",
    "        2 ** trial.suggest_int(name=\"batch_size_exponent\", low=7, high=7)\n",
    "    )\n",
    "    experiment.log_parameter(name=\"batch_size\", value=batch_size)\n",
    "    train_iter, train_len, dev_iter, dev_len = get_train_dev_iterators(\n",
    "        dataset=dataset, first_size=int(len(dataset) * 0.95), batch_size=batch_size\n",
    "    )\n",
    "    experiment.log_parameters(\n",
    "        dic={\"train_set_samples\": train_len, \"dev_set_samples\": dev_len}\n",
    "    )\n",
    "\n",
    "    ## Compile Model\n",
    "    learning_rate: float = trial.suggest_discrete_uniform(\n",
    "        name=\"learning_rate\", high=2.0e-4, low=1.0e-4, q=0.1e-4\n",
    "    )\n",
    "    if resume_experiment_key is None:\n",
    "        num_residual_blocks: int = trial.suggest_int(\n",
    "            name=\"num_residual_blocks\", low=12, high=12\n",
    "        )\n",
    "        residual_scaling: float = trial.suggest_discrete_uniform(\n",
    "            name=\"residual_scaling\", low=0.1, high=0.3, q=0.05\n",
    "        )\n",
    "    else:  # resume training from a previous experiment\n",
    "        for _model_type in (\n",
    "            [\"generator_model\", \"discriminator_model\"]\n",
    "            if reload_d_model_weights\n",
    "            else [\"generator_model\"]\n",
    "        ):\n",
    "            hyperparameters = _download_model_weights_from_comet(\n",
    "                experiment_key=resume_experiment_key,\n",
    "                download_path=f\"{base_model_weight_path}/srgan_{_model_type}_weights.npz\",\n",
    "            )\n",
    "        num_residual_blocks = int(hyperparameters[\"num_residual_blocks\"])\n",
    "        residual_scaling = float(hyperparameters[\"residual_scaling\"])\n",
    "        # learning_rate = float(hyperparameters[\"generator_lr\"])\n",
    "\n",
    "    g_model, g_optimizer, d_model, d_optimizer = compile_srgan_model(\n",
    "        num_residual_blocks=num_residual_blocks,\n",
    "        residual_scaling=residual_scaling,\n",
    "        learning_rate=learning_rate,\n",
    "    )\n",
    "    if resume_experiment_key is not None:\n",
    "        chainer.serializers.load_npz(\n",
    "            file=f\"{base_model_weight_path}/srgan_generator_model_weights.npz\",\n",
    "            obj=g_model,\n",
    "        )\n",
    "        if reload_d_model_weights:\n",
    "            chainer.serializers.load_npz(\n",
    "                file=f\"{base_model_weight_path}/srgan_discriminator_model_weights.npz\",\n",
    "                obj=d_model,\n",
    "            )\n",
    "    experiment.log_parameters(\n",
    "        dic={\n",
    "            \"num_residual_blocks\": g_model.num_residual_blocks,\n",
    "            \"residual_scaling\": g_model.residual_scaling,\n",
    "            \"generator_optimizer\": \"adam\",\n",
    "            \"generator_lr\": g_optimizer.alpha,  # learning rate\n",
    "            \"generator_epsilon\": g_optimizer.eps,  # epsilon\n",
    "            \"discriminator_optimizer\": \"adam\",\n",
    "            \"discriminator_lr\": d_optimizer.alpha,  # learning rate\n",
    "            \"discriminator_adam_epsilon\": d_optimizer.eps,  # epsilon\n",
    "        }\n",
    "    )\n",
    "\n",
    "    ## Run Trainer and save trained model\n",
    "    epochs: int = trial.suggest_int(name=\"num_epochs\", low=15, high=150)\n",
    "    experiment.log_parameter(name=\"num_epochs\", value=epochs)\n",
    "\n",
    "    metric_names = [\n",
    "        \"discriminator_loss\",\n",
    "        \"discriminator_accu\",\n",
    "        \"generator_loss\",\n",
    "        \"generator_psnr\",\n",
    "        \"generator_ssim\",\n",
    "    ]\n",
    "    columns = metric_names + [f\"val_{metric_name}\" for metric_name in metric_names]\n",
    "    dataframe = pd.DataFrame(index=np.arange(epochs), columns=columns)\n",
    "    progressbar = tqdm.tqdm(unit=\"epoch\", total=epochs, position=0)\n",
    "\n",
    "    train_iter.reset()\n",
    "    dev_iter.reset()\n",
    "\n",
    "    base_rmse = 500  # will save and upload models that beat this score\n",
    "    best_rmse_test = 250\n",
    "    for i in range(epochs):\n",
    "        metrics_dict = trainer(\n",
    "            i=i,\n",
    "            columns=columns,\n",
    "            train_iter=train_iter,\n",
    "            dev_iter=dev_iter,\n",
    "            g_model=g_model,\n",
    "            g_optimizer=g_optimizer,\n",
    "            d_model=d_model,\n",
    "            d_optimizer=d_optimizer,\n",
    "        )\n",
    "\n",
    "        ## Record training loss and metric info, and plot using livelossplot if enabled\n",
    "        dataframe.loc[i] = [\n",
    "            np.mean(metrics_dict[metric]) for metric in dataframe.keys()\n",
    "        ]\n",
    "        epoch_metrics = dataframe.loc[i].to_dict()\n",
    "        if enable_livelossplot == True:\n",
    "            livelossplot.draw_plot(\n",
    "                logs=dataframe.to_dict(orient=\"records\"),\n",
    "                metrics=metric_names,\n",
    "                max_cols=5,\n",
    "                figsize=(21, 9),\n",
    "                max_epoch=None,\n",
    "            )\n",
    "        progressbar.set_postfix(ordered_dict=epoch_metrics)\n",
    "        progressbar.update(n=1)\n",
    "        experiment.log_metrics(dic=epoch_metrics, step=i)\n",
    "\n",
    "        ## Evaluate model and produce image of test area\n",
    "        rmse_test, predicted_test_grid = get_deepbedmap_test_result(model=g_model)\n",
    "        experiment.log_metric(name=\"rmse_test\", value=rmse_test)\n",
    "        fig = gmt.Figure()\n",
    "        fig.grdimage(\n",
    "            grid=predicted_test_grid,\n",
    "            projection=\"x1:2000000\",\n",
    "            frame=[\"WSne\", \"af\"],\n",
    "            cmap=\"oleron\",\n",
    "        )\n",
    "        fig.colorbar(S=True, position=\"JMR+n\", frame=\"af\")\n",
    "        fig.savefig(f\"model/images/{experiment.get_key()}.png\")\n",
    "        experiment.log_image(\n",
    "            name=\"predicted_test_grid\",\n",
    "            # image_data=np.flipud(predicted_test_grid.data),\n",
    "            image_data=f\"model/images/{experiment.get_key()}.png\",\n",
    "            # image_colormap=\"BrBG\",\n",
    "        )\n",
    "        trial.report(value=rmse_test, step=i)\n",
    "\n",
    "        # Save generator and discriminator neural network model weights,\n",
    "        # and save only generator model architecture\n",
    "        if rmse_test < best_rmse_test:\n",
    "            best_rmse_test = rmse_test\n",
    "            (\n",
    "                g_model_weights_path,\n",
    "                d_model_weights_path,\n",
    "                g_model_architecture_path,\n",
    "            ) = save_model_weights_and_architecture(\n",
    "                generator_model=g_model,\n",
    "                discriminator_model=d_model,\n",
    "                save_path=base_model_weight_path,\n",
    "            )\n",
    "\n",
    "        # Upload neural network weights if this is the final epoch,\n",
    "        # or if the trial should be pruned and we have a good RMSE_test result\n",
    "        if best_rmse_test < base_rmse and (i == epochs - 1 or trial.should_prune()):\n",
    "            experiment.log_asset(\n",
    "                file_data=g_model_weights_path,\n",
    "                file_name=os.path.basename(g_model_weights_path),\n",
    "            )\n",
    "            experiment.log_asset(\n",
    "                file_data=d_model_weights_path,\n",
    "                file_name=os.path.basename(d_model_weights_path),\n",
    "            )\n",
    "            try:\n",
    "                with open(file=g_model_architecture_path) as outgraph:\n",
    "                    experiment.set_model_graph(graph=outgraph.read())\n",
    "                experiment.log_asset(\n",
    "                    file_data=g_model_architecture_path,\n",
    "                    file_name=os.path.basename(g_model_architecture_path),\n",
    "                )\n",
    "            except FileNotFoundError:\n",
    "                print(f\"Could not find {g_model_architecture_path}\")\n",
    "            for f in glob.glob(f\"{base_model_weight_path}/*\"):\n",
    "                shutil.copy2(src=f, dst=\"model/weights\")\n",
    "            try:\n",
    "                shutil.rmtree(path=base_model_weight_path)\n",
    "            except FileNotFoundError:\n",
    "                pass\n",
    "\n",
    "        ## Pruning unpromising trials with vanishing/exploding gradients\n",
    "        if (\n",
    "            trial.should_prune()\n",
    "            or epoch_metrics[\"generator_psnr\"] < 0\n",
    "            or np.isnan(epoch_metrics[\"generator_loss\"])\n",
    "            or np.isnan(epoch_metrics[\"discriminator_loss\"])\n",
    "        ):\n",
    "            experiment.end()\n",
    "            raise optuna.structs.TrialPruned()\n",
    "\n",
    "    ## Upload final predicted figure with scalebar\n",
    "    predicted_test_grid.plot.imshow(\n",
    "        cmap=\"BrBG\",\n",
    "        size=14,\n",
    "        aspect=predicted_test_grid.shape[1] / predicted_test_grid.shape[0],\n",
    "    )\n",
    "    plt.tight_layout()\n",
    "    experiment.log_figure(figure_name=\"final_figure_of_predicted_test_area\")\n",
    "    plt.close()\n",
    "\n",
    "    print(f\"Experiment yielded Root Mean Square Error of {rmse_test:.2f} on test set\")\n",
    "    experiment.end()\n",
    "\n",
    "    return rmse_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "n_trials = 90\n",
    "if n_trials == 1:  # run training once only, i.e. just test the objective function\n",
    "    objective(enable_livelossplot=True, enable_comet_logging=True)\n",
    "elif n_trials > 1:  # perform hyperparameter tuning with multiple experimental trials\n",
    "    # Set different seed using len($HOSTNAME) + GPU_ID\n",
    "    hostname: str = os.uname().nodename\n",
    "    tpe_seed = len(hostname) + int(os.getenv(key=\"CUDA_VISIBLE_DEVICES\", default=\"0\"))\n",
    "    # Tree-structured Parzen Estimator using HyperOpt defaults\n",
    "    sampler = optuna.samplers.TPESampler(\n",
    "        seed=tpe_seed, **optuna.samplers.TPESampler.hyperopt_parameters()\n",
    "    )\n",
    "    study = optuna.create_study(\n",
    "        study_name=\"DeepBedMap_tuning\",\n",
    "        storage=f\"sqlite:///model/logs/train_on_{hostname}.db\",\n",
    "        sampler=sampler,\n",
    "        pruner=optuna.pruners.HyperbandPruner(\n",
    "            min_resource=15,  # minimum number of epochs (i.e. warmup period)\n",
    "            max_resource=150,  # maximum number of epochs\n",
    "            reduction_factor=3,\n",
    "        ),\n",
    "        load_if_exists=True,\n",
    "    )\n",
    "    study.optimize(func=objective, n_trials=n_trials, n_jobs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>number</th>\n",
       "      <th>value</th>\n",
       "      <th>datetime_start</th>\n",
       "      <th>datetime_complete</th>\n",
       "      <th>duration</th>\n",
       "      <th>params_batch_size_exponent</th>\n",
       "      <th>params_learning_rate</th>\n",
       "      <th>params_num_epochs</th>\n",
       "      <th>system_attrs_completed_rung_0</th>\n",
       "      <th>system_attrs_completed_rung_1</th>\n",
       "      <th>system_attrs_completed_rung_2</th>\n",
       "      <th>state</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>32</td>\n",
       "      <td>195.119021</td>\n",
       "      <td>2020-08-09 06:39:19.698288</td>\n",
       "      <td>2020-08-09 07:40:30.935322</td>\n",
       "      <td>01:01:11.237034</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00018</td>\n",
       "      <td>139</td>\n",
       "      <td>219.913707</td>\n",
       "      <td>229.054037</td>\n",
       "      <td>196.136216</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>26</td>\n",
       "      <td>196.244989</td>\n",
       "      <td>2020-08-09 04:48:53.181358</td>\n",
       "      <td>2020-08-09 05:52:27.865477</td>\n",
       "      <td>01:03:34.684119</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00017</td>\n",
       "      <td>138</td>\n",
       "      <td>199.795300</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>33</td>\n",
       "      <td>196.384394</td>\n",
       "      <td>2020-08-09 06:56:14.155940</td>\n",
       "      <td>2020-08-09 07:47:57.408358</td>\n",
       "      <td>00:51:43.252418</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00016</td>\n",
       "      <td>139</td>\n",
       "      <td>225.771761</td>\n",
       "      <td>225.379145</td>\n",
       "      <td>202.567884</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>63</th>\n",
       "      <td>63</td>\n",
       "      <td>196.778727</td>\n",
       "      <td>2020-08-09 14:05:53.858627</td>\n",
       "      <td>2020-08-09 15:15:09.993879</td>\n",
       "      <td>01:09:16.135252</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00016</td>\n",
       "      <td>142</td>\n",
       "      <td>220.912740</td>\n",
       "      <td>194.858577</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>198.460686</td>\n",
       "      <td>2020-08-08 22:52:10.427130</td>\n",
       "      <td>2020-08-08 23:58:05.890701</td>\n",
       "      <td>01:05:55.463571</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00017</td>\n",
       "      <td>135</td>\n",
       "      <td>226.985104</td>\n",
       "      <td>212.918729</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>35</td>\n",
       "      <td>199.122106</td>\n",
       "      <td>2020-08-09 07:40:31.104526</td>\n",
       "      <td>2020-08-09 08:29:55.166656</td>\n",
       "      <td>00:49:24.062130</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00018</td>\n",
       "      <td>105</td>\n",
       "      <td>222.343048</td>\n",
       "      <td>223.767694</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>30</td>\n",
       "      <td>199.773410</td>\n",
       "      <td>2020-08-09 05:52:27.991434</td>\n",
       "      <td>2020-08-09 06:59:37.321033</td>\n",
       "      <td>01:07:09.329599</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00017</td>\n",
       "      <td>139</td>\n",
       "      <td>233.206111</td>\n",
       "      <td>233.286785</td>\n",
       "      <td>211.422163</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50</th>\n",
       "      <td>50</td>\n",
       "      <td>200.702349</td>\n",
       "      <td>2020-08-09 11:07:08.013653</td>\n",
       "      <td>2020-08-09 12:06:19.533297</td>\n",
       "      <td>00:59:11.519644</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00019</td>\n",
       "      <td>133</td>\n",
       "      <td>210.225846</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>24</td>\n",
       "      <td>200.889667</td>\n",
       "      <td>2020-08-09 03:49:05.262452</td>\n",
       "      <td>2020-08-09 04:50:31.636901</td>\n",
       "      <td>01:01:26.374449</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00015</td>\n",
       "      <td>126</td>\n",
       "      <td>224.702599</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>25</td>\n",
       "      <td>202.500655</td>\n",
       "      <td>2020-08-09 04:27:56.494916</td>\n",
       "      <td>2020-08-09 05:16:59.428692</td>\n",
       "      <td>00:49:02.933776</td>\n",
       "      <td>7</td>\n",
       "      <td>0.00016</td>\n",
       "      <td>132</td>\n",
       "      <td>242.006534</td>\n",
       "      <td>224.059256</td>\n",
       "      <td>NaN</td>\n",
       "      <td>COMPLETE</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    number       value             datetime_start          datetime_complete  \\\n",
       "32      32  195.119021 2020-08-09 06:39:19.698288 2020-08-09 07:40:30.935322   \n",
       "26      26  196.244989 2020-08-09 04:48:53.181358 2020-08-09 05:52:27.865477   \n",
       "33      33  196.384394 2020-08-09 06:56:14.155940 2020-08-09 07:47:57.408358   \n",
       "63      63  196.778727 2020-08-09 14:05:53.858627 2020-08-09 15:15:09.993879   \n",
       "2        2  198.460686 2020-08-08 22:52:10.427130 2020-08-08 23:58:05.890701   \n",
       "35      35  199.122106 2020-08-09 07:40:31.104526 2020-08-09 08:29:55.166656   \n",
       "30      30  199.773410 2020-08-09 05:52:27.991434 2020-08-09 06:59:37.321033   \n",
       "50      50  200.702349 2020-08-09 11:07:08.013653 2020-08-09 12:06:19.533297   \n",
       "24      24  200.889667 2020-08-09 03:49:05.262452 2020-08-09 04:50:31.636901   \n",
       "25      25  202.500655 2020-08-09 04:27:56.494916 2020-08-09 05:16:59.428692   \n",
       "\n",
       "          duration  params_batch_size_exponent  params_learning_rate  \\\n",
       "32 01:01:11.237034                           7               0.00018   \n",
       "26 01:03:34.684119                           7               0.00017   \n",
       "33 00:51:43.252418                           7               0.00016   \n",
       "63 01:09:16.135252                           7               0.00016   \n",
       "2  01:05:55.463571                           7               0.00017   \n",
       "35 00:49:24.062130                           7               0.00018   \n",
       "30 01:07:09.329599                           7               0.00017   \n",
       "50 00:59:11.519644                           7               0.00019   \n",
       "24 01:01:26.374449                           7               0.00015   \n",
       "25 00:49:02.933776                           7               0.00016   \n",
       "\n",
       "    params_num_epochs  system_attrs_completed_rung_0  \\\n",
       "32                139                     219.913707   \n",
       "26                138                     199.795300   \n",
       "33                139                     225.771761   \n",
       "63                142                     220.912740   \n",
       "2                 135                     226.985104   \n",
       "35                105                     222.343048   \n",
       "30                139                     233.206111   \n",
       "50                133                     210.225846   \n",
       "24                126                     224.702599   \n",
       "25                132                     242.006534   \n",
       "\n",
       "    system_attrs_completed_rung_1  system_attrs_completed_rung_2     state  \n",
       "32                     229.054037                     196.136216  COMPLETE  \n",
       "26                            NaN                            NaN  COMPLETE  \n",
       "33                     225.379145                     202.567884  COMPLETE  \n",
       "63                     194.858577                            NaN  COMPLETE  \n",
       "2                      212.918729                            NaN  COMPLETE  \n",
       "35                     223.767694                            NaN  COMPLETE  \n",
       "30                     233.286785                     211.422163  COMPLETE  \n",
       "50                            NaN                            NaN  COMPLETE  \n",
       "24                            NaN                            NaN  COMPLETE  \n",
       "25                     224.059256                            NaN  COMPLETE  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if n_trials > 1:\n",
    "    study = optuna.load_study(\n",
    "        study_name=\"DeepBedMap_tuning\",\n",
    "        storage=f\"sqlite:///model/logs/train_on_{hostname}.db\",\n",
    "    )\n",
    "    topten_df = study.trials_dataframe().nsmallest(n=10, columns=\"value\")\n",
    "    IPython.display.display(topten_df)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "deepbedmap",
   "language": "python",
   "name": "deepbedmap"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}