DNN_models_comparison.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d545b75d",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "### Import necessary modules and do some basic setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1427bd4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python ≥3.5 is required\n",
    "import sys\n",
    "assert sys.version_info >= (3, 5)\n",
    "\n",
    "# Scikit-Learn ≥0.20 is required\n",
    "import sklearn\n",
    "assert sklearn.__version__ >= '0.20'\n",
    "\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import confusion_matrix, classification_report, ConfusionMatrixDisplay, precision_score, recall_score, roc_auc_score, roc_curve\n",
    "from sklearn.utils import class_weight\n",
    "\n",
    "# TensorFlow ≥2.0 is required\n",
    "import tensorflow_addons as tfa\n",
    "import tensorflow as tf\n",
    "assert tf.__version__ >= '2.0'\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers, regularizers\n",
    "\n",
    "print(\"Num GPUs Available: \", len(tf.config.list_physical_devices('GPU')))\n",
    "\n",
    "# Common imports\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import xarray as xr\n",
    "import dask\n",
    "import datetime\n",
    "import math\n",
    "import pickle\n",
    "import pathlib\n",
    "import hashlib\n",
    "dask.config.set({'array.slicing.split_large_chunks': False})\n",
    "\n",
    "# To make this notebook's output stable across runs\n",
    "np.random.seed(42)\n",
    "\n",
    "# Config matplotlib\n",
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "mpl.rc('axes', labelsize=14)\n",
    "mpl.rc('xtick', labelsize=12)\n",
    "mpl.rc('ytick', labelsize=12)\n",
    "\n",
    "# Dotenv\n",
    "from dotenv import dotenv_values\n",
    "\n",
    "# Custom utils\n",
    "from utils.utils_data import *\n",
    "from utils.utils_ml import *\n",
    "from utils.utils_resnet import *\n",
    "from utils.utils_plot import *\n",
    "from utils.utils_unet import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41766d4a",
   "metadata": {},
   "source": [
    "### Define some paths and constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc1d02d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = dotenv_values(\".env\")\n",
    "\n",
    "# Paths\n",
    "PATH_ERA5 = config['PATH_ERA5']\n",
    "PATH_EOBS = config['PATH_EOBS']\n",
    "\n",
    "# Some constants\n",
    "G = 9.80665\n",
    "\n",
    "# Options\n",
    "DATE_START = '1979-01-01'\n",
    "DATE_END = '2021-12-31'\n",
    "YY_TRAIN = [1979, 2015]\n",
    "YY_TEST = [2016, 2021]\n",
    "LEVELS = [300, 500, 700, 850, 925, 1000]\n",
    "LONS_INPUT = [-25, 30]\n",
    "LATS_INPUT = [30, 75]\n",
    "LONS_PREC = [-25, 30]\n",
    "LATS_PREC = [30, 75]\n",
    "BATCH_SIZE = 64\n",
    "PRECIP_DATA = 'ERA5-low' # Options: ERA5-hi, ERA5-low, E-OBS\n",
    "PRECIP_XTRM = 0.95 # Percentile (threshold) for the extremes\n",
    "USE_3D_ONLY = False # False\n",
    "CREATE_MASK_EOBS = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c3dbae",
   "metadata": {},
   "source": [
    "# Data preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74596028",
   "metadata": {},
   "source": [
    "## Target variable: precipitation field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0abeea2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load precipitation\n",
    "if PRECIP_DATA == 'ERA5-hi':\n",
    "    pr = get_nc_data(PATH_ERA5 + '/precipitation/orig_grid/daily/*nc', DATE_START, DATE_END, LONS_PREC, LATS_PREC)\n",
    "    pr = pr.tp\n",
    "elif PRECIP_DATA == 'ERA5-low':\n",
    "    pr = get_nc_data(PATH_ERA5 + '/precipitation/day_grid1/*nc', DATE_START, DATE_END, LONS_PREC, LATS_PREC)\n",
    "    pr = pr.tp\n",
    "elif PRECIP_DATA in ['E-OBS', 'EOBS']:\n",
    "    pr = get_nc_data(PATH_EOBS + '/eobs_1deg_v26.0e.nc', DATE_START, DATE_END, LONS_PREC, LATS_PREC)\n",
    "    pr = pr.rr\n",
    "    CREATE_MASK_EOBS = True\n",
    "else:\n",
    "    raise ValueError('Precipitation data not well defined')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1344518f-75f4-4227-9c98-e5e7261acc69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Invert lat axis if needed\n",
    "if pr.lat[0].values < pr.lat[1].values:\n",
    "    pr = pr.reindex(lat=list(reversed(pr.lat)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f56ebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a dimension to be used as channel in the DNN\n",
    "pr = pr.expand_dims('level', -1)\n",
    "pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62223e8f-9b85-48f7-91a7-3693d79e93d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mask\n",
    "mask = None\n",
    "if CREATE_MASK_EOBS:\n",
    "    if PRECIP_DATA in ['E-OBS', 'EOBS']:\n",
    "        peobs = pr[:,:,:,0]\n",
    "    else:\n",
    "        peobs = get_nc_data(PATH_EOBS + '/eobs_1deg_v26.0e.nc', DATE_START, DATE_END, LONS_PREC, LATS_PREC)\n",
    "        peobs = peobs.rr\n",
    "        if peobs.lat[0].values < peobs.lat[1].values:\n",
    "            peobs = peobs.reindex(lat=list(reversed(peobs.lat)))\n",
    "    mask = np.isnan(peobs[0,:,:])\n",
    "    mask = np.invert(mask)\n",
    "    mask.plot()\n",
    "    mask = mask.to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67518341-5e50-4db4-8e4b-c9b009977ba6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mask out ERA5 precip\n",
    "if CREATE_MASK_EOBS and PRECIP_DATA in ['ERA5-low', 'ERA5-hi']:\n",
    "    print('Applying mask to the data')\n",
    "    nan_mask = mask*1.0\n",
    "    nan_mask[nan_mask == 0] = np.nan\n",
    "    pr = pr[:,:,:,0]*nan_mask\n",
    "    pr = pr.expand_dims('level', -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a9308b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the extreme exceedence\n",
    "qq = xr.DataArray(pr).chunk(dict(time=-1)).quantile(PRECIP_XTRM, dim='time')\n",
    "pr_xtrm = xr.DataArray(pr > qq)\n",
    "pr_xtrm = pr_xtrm*1 # Transform to number"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43d2a68-8e48-4a35-a7df-45352f67640d",
   "metadata": {},
   "outputs": [],
   "source": [
    "qq.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe1f42e-d1dc-4b77-b399-2a0ed5757e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the values to plot later\n",
    "if PRECIP_XTRM == 0.95:\n",
    "    np.save('tmp/data/qq_095.npy',qq.to_numpy().squeeze())\n",
    "elif PRECIP_XTRM == 0.99:\n",
    "    np.save('tmp/data/qq_099.npy',qq.to_numpy().squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd208aba-31eb-42f2-a1b0-6d643b2ceadc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract coordinates for precip\n",
    "lats_y = pr.lat.to_numpy()\n",
    "lons_y = pr.lon.to_numpy()\n",
    "np.save('tmp/data/lats_y.npy',lats_y.squeeze())\n",
    "np.save('tmp/data/lons_y.npy',lons_y.squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "175c2dbe-bd1d-49a3-b94b-5f949bc04bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the date with the max number of grid points with pr > threshold\n",
    "x = pr_xtrm.sum(dim=['lat', 'lon'])\n",
    "plot_date = x.argmax(dim='time')\n",
    "\n",
    "# Plot precip, threshold exceedance and threshold values\n",
    "fig, axs = plt.subplots(1, 3, figsize=(21,4))\n",
    "\n",
    "pp = pr[plot_date,:,:,0]\n",
    "px = pr_xtrm[plot_date,:,:,0]\n",
    "\n",
    "plot_map(axs[0], lons_y, lats_y, pp.to_numpy().squeeze(), title=\"Daily precipitation\", cmap=mpl.cm.YlGnBu)\n",
    "plot_map(axs[1], lons_y, lats_y, px.to_numpy().squeeze(), title=\"Threshold exceedance\", cmap=mpl.cm.YlGnBu)\n",
    "plot_map(axs[2], lons_y, lats_y, qq.to_numpy().squeeze(), title=\"Threshold values\", cmap=mpl.cm.YlGnBu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c847649a",
   "metadata": {},
   "source": [
    "## Input data: meteorological fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce5176d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load geopotential height\n",
    "z = get_era5_data(PATH_ERA5 + '/geopotential/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "z['time'] = pd.DatetimeIndex(z.time.dt.date)\n",
    "z = z.sel(level=LEVELS)\n",
    "\n",
    "# Get Z in geopotential height (m)\n",
    "z.z.values = z.z.values/G\n",
    "\n",
    "# Load temperature\n",
    "t = get_era5_data(PATH_ERA5 + '/temperature/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "t['time'] = pd.DatetimeIndex(t.time.dt.date)\n",
    "\n",
    "# Load relative humidity\n",
    "rh = get_era5_data(PATH_ERA5 + '/relative_humidity/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "rh['time'] = pd.DatetimeIndex(rh.time.dt.date)\n",
    "rh = rh.sel(level=LEVELS)\n",
    "\n",
    "# Load total column water\n",
    "if not USE_3D_ONLY:\n",
    "    tcw = get_era5_data(PATH_ERA5 + '/total_column_water/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "    tcw['time'] = pd.DatetimeIndex(tcw.time.dt.date)\n",
    "\n",
    "# Load wind components\n",
    "u = get_era5_data(PATH_ERA5 + '/U_wind/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "u['time'] = pd.DatetimeIndex(u.time.dt.date)\n",
    "v = get_era5_data(PATH_ERA5 + '/V_wind/day_grid1/*.nc', DATE_START, DATE_END, LONS_INPUT, LATS_INPUT)\n",
    "v['time'] = pd.DatetimeIndex(v.time.dt.date)\n",
    "\n",
    "# Checking dimensions\n",
    "print('dimension of pr:', pr.dims)\n",
    "print('dimension of z', z.dims)\n",
    "print('dimension of t:', t.dims)\n",
    "print('dimension of rh:', rh.dims)\n",
    "if not USE_3D_ONLY:\n",
    "    print('dimension of tcw:', tcw.dims)\n",
    "print('dimension of u:', u.dims)\n",
    "print('dimension of v:', v.dims)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69d54b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge arrays\n",
    "if USE_3D_ONLY:\n",
    "    X = xr.merge([z, t, rh, u, v])\n",
    "else:\n",
    "    X = xr.merge([z, t, rh, tcw, u, v])\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a200bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Invert lat axis if needed\n",
    "if X.lat[0].values < X.lat[1].values:\n",
    "    X = X.reindex(lat=list(reversed(X.lat)))\n",
    "    \n",
    "# Get axes\n",
    "lats_x = X.lat\n",
    "lons_x = X.lon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de7643ef",
   "metadata": {},
   "source": [
    "### Split data and transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23738ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split into training and test\n",
    "X_train_full = X.sel(time=slice(f'{YY_TRAIN[0]}-01-01', f'{YY_TRAIN[1]}-12-31'))\n",
    "X_test = X.sel(time=slice(f'{YY_TEST[0]}-01-01', f'{YY_TEST[1]}-12-31'))\n",
    "\n",
    "pr_train_full = pr.sel(time=slice(f'{YY_TRAIN[0]}-01-01', f'{YY_TRAIN[1]}-12-31'))\n",
    "pr_test = pr.sel(time=slice(f'{YY_TEST[0]}-01-01', f'{YY_TEST[1]}-12-31'))\n",
    "\n",
    "pr_xtrm_train_full = pr_xtrm.sel(time=slice(f'{YY_TRAIN[0]}-01-01', f'{YY_TRAIN[1]}-12-31'))\n",
    "pr_xtrm_test = pr_xtrm.sel(time=slice(f'{YY_TEST[0]}-01-01', f'{YY_TEST[1]}-12-31'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be5e320",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a data generator\n",
    "if USE_3D_ONLY:\n",
    "    dic = {'z': LEVELS,\n",
    "       't': LEVELS,\n",
    "       'r': LEVELS,\n",
    "       'u': LEVELS,\n",
    "       'v': LEVELS}\n",
    "else:\n",
    "    dic = {'z': LEVELS,\n",
    "       't': LEVELS,\n",
    "       'r': LEVELS,\n",
    "       'tcwv': None,\n",
    "       'u': LEVELS,\n",
    "       'v': LEVELS}\n",
    "\n",
    "from utils.utils_ml import *\n",
    "\n",
    "YY_VALID = 2005\n",
    "\n",
    "dg_train = DataGeneratorWithExtremes(X_train_full.sel(time=slice(f'{YY_TRAIN[0]}', f'{YY_VALID}')),\n",
    "                                     pr_train_full.sel(time=slice(f'{YY_TRAIN[0]}', f'{YY_VALID}')),\n",
    "                                     pr_xtrm_train_full.sel(time=slice(f'{YY_TRAIN[0]}', f'{YY_VALID}')),\n",
    "                                     dic, batch_size=BATCH_SIZE, load=True)\n",
    "dg_valid = DataGeneratorWithExtremes(X_train_full.sel(time=slice(f'{YY_VALID+1}', f'{YY_TRAIN[1]}')),\n",
    "                                     pr_train_full.sel(time=slice(f'{YY_VALID+1}', f'{YY_TRAIN[1]}')),\n",
    "                                     pr_xtrm_train_full.sel(time=slice(f'{YY_VALID+1}', f'{YY_TRAIN[1]}')),\n",
    "                                     dic, mean=dg_train.mean, std=dg_train.std,\n",
    "                                     batch_size=BATCH_SIZE, load=True)\n",
    "dg_test = DataGeneratorWithExtremes(X_test, pr_test, pr_xtrm_test, dic,\n",
    "                                    mean=dg_train.mean, std=dg_train.std,\n",
    "                                    batch_size=BATCH_SIZE, load=True, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "429afe9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "i_shape = dg_train.X.shape[1:]\n",
    "o_shape = dg_train.y.shape[1:]\n",
    "\n",
    "print(f'X shape: {i_shape}')\n",
    "print(f'y shape: {o_shape}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f2c3027-590c-4b4b-9fb5-ea536220620c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save precip test data\n",
    "y_xtrm = dg_test.y_xtrm.to_numpy().squeeze()\n",
    "y_prec = dg_test.y.to_numpy().squeeze()\n",
    "np.save('tmp/data/y_prec.npy',y_prec)\n",
    "if PRECIP_XTRM == 0.95:\n",
    "    np.save('tmp/data/y_xtrm_095.npy', y_xtrm)\n",
    "elif PRECIP_XTRM == 0.99:\n",
    "    np.save('tmp/data/y_xtrm_099.npy', y_xtrm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8d1475-9b93-4ed7-8fd8-a78894d7e62b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the mean precipitation\n",
    "fig = plt.figure(figsize=(7,5))\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "vals = np.mean(dg_train.y, axis=0).squeeze()\n",
    "plot_map(ax, lons_y, lats_y, vals, title=\"Mean precipitation\", cmap=mpl.cm.YlGnBu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afb29c12-e9c3-4a97-b0c1-88ad4ff1de7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the mean of the predictors\n",
    "n_figs = len(dg_train.X[0,0,0,:])\n",
    "ncols = 5\n",
    "nrows = -(-n_figs // ncols)\n",
    "fig, axes = plt.subplots(figsize=(24, 3.3*nrows), ncols=ncols, nrows=nrows)\n",
    "for i in range(n_figs):\n",
    "    i_row = i // ncols\n",
    "    i_col = i % ncols\n",
    "    ax = axes[i_row, i_col]\n",
    "    vals = np.mean(dg_train.X[:,:,:,i], axis=0).to_numpy()\n",
    "    plot_map(ax, lons_x, lats_x, vals, title=f\"Average of feature {i+1}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ba7bd3",
   "metadata": {},
   "source": [
    "# Model creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4bb749",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepFactory(tf.keras.Model):\n",
    "    \"\"\"\n",
    "    Model factory.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, arch, input_size, output_size, for_extremes=False, latent_dim=128, \n",
    "                 dropout_rate=0.2, use_batch_norm=True, inner_activation='relu', unet_filters_nb=64, \n",
    "                 unet_depth=4, use_upsample=False, output_scaling=1, output_crop=None, levels_nb_3d=6,\n",
    "                 mask=None):\n",
    "        super(DeepFactory, self).__init__()\n",
    "        self.arch = arch\n",
    "        self.input_size = list(input_size)\n",
    "        self.output_size = list(output_size)\n",
    "        self.for_extremes = for_extremes\n",
    "        self.latent_dim = latent_dim\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.use_batch_norm = use_batch_norm\n",
    "        self.use_upsample = use_upsample\n",
    "        self.inner_activation = inner_activation\n",
    "        self.unet_depth = unet_depth\n",
    "        self.unet_filters_nb = unet_filters_nb\n",
    "        self.output_scaling = output_scaling\n",
    "        self.output_crop = output_crop\n",
    "        self.levels_nb_3d = levels_nb_3d\n",
    "        self.mask = mask\n",
    "        \n",
    "        self.last_activation = 'relu'\n",
    "        if for_extremes:\n",
    "            self.last_activation = 'sigmoid'\n",
    "\n",
    "        if arch == 'Davenport-2021':\n",
    "            self.build_Davenport_2021()\n",
    "        elif arch == 'CNN-2L':\n",
    "            self.build_CNN_2L()\n",
    "        elif arch == 'Unet':\n",
    "            self.build_Unet()\n",
    "        elif arch == 'Unet-attention':\n",
    "            self.build_UnetAtt()\n",
    "        elif arch == 'Pan-2019':\n",
    "            self.build_Pan_2019()\n",
    "        elif arch =='Conv-LTSM':\n",
    "            self.build_convLTSM()\n",
    "        elif arch =='RaNet':\n",
    "            self.build_RaNet()\n",
    "        else:\n",
    "            raise ValueError('The architecture was not correctly defined')\n",
    "        \n",
    "        \n",
    "    def build_Davenport_2021(self):\n",
    "        \"\"\"\n",
    "        Based on: Davenport, F. V., & Diffenbaugh, N. S. (2021). Using Machine Learning to \n",
    "        Analyze Physical Causes of Climate Change: A Case Study of U.S. Midwest Extreme Precipitation. \n",
    "        Geophysical Research Letters, 48(15). https://doi.org/10.1029/2021GL093787\n",
    "        \"\"\"\n",
    "        \n",
    "        # Downsampling\n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        x = layers.Conv2D(16, 3, padding='same', activity_regularizer=regularizers.l2(0.01))(inputs)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.MaxPooling2D(pool_size=2)(x)\n",
    "        x = layers.SpatialDropout2D(self.dropout_rate)(x) # In original: simple Dropout\n",
    "        x = layers.Conv2D(16, 3, padding='same', activity_regularizer=regularizers.l2(0.01))(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.MaxPooling2D(pool_size=2)(x)\n",
    "        x = layers.SpatialDropout2D(self.dropout_rate)(x) # In original: simple Dropout\n",
    "        x = layers.Flatten()(x)                \n",
    "        x = layers.Dense(self.latent_dim, activity_regularizer=regularizers.l2(0.001))(x) # In original: 16\n",
    "        x = layers.Activation('relu')(x)\n",
    "\n",
    "        next_shape = self.get_shape_for(stride_factor=4)\n",
    "\n",
    "        # Upsampling. In original: no decoder\n",
    "        x = self.dense_block(x, np.prod(next_shape))\n",
    "        x = layers.Reshape(target_shape=next_shape)(x)\n",
    "        x = self.deconv_block(x, 16, 3, stride=2)\n",
    "        x = self.deconv_block(x, 16, 3, stride=2)\n",
    "        x = self.conv_block(x, 1, 3, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    " \n",
    "        self.model = keras.Model(inputs, outputs, name=\"Davenport-2021\")\n",
    "        \n",
    "        \n",
    "    def build_Pan_2019(self):\n",
    "        \"\"\"\n",
    "        Based on: Pan, B., Hsu, K., AghaKouchak, A., & Sorooshian, S. (2019). \n",
    "        Improving Precipitation Estimation Using Convolutional Neural Network. \n",
    "        Water Resources Research, 55(3), 2301–2321. https://doi.org/10.1029/2018WR024090\n",
    "        \"\"\"\n",
    "        # In original: padding='valid'\n",
    "        \n",
    "        # Downsampling\n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        x = layers.Conv2D(15, 4, padding='same')(inputs)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.Conv2D(20, 4, padding='same')(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.MaxPooling2D(pool_size=2)(x)\n",
    "        x = layers.Conv2D(20, 4, padding='same')(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.MaxPooling2D(pool_size=2)(x)\n",
    "        x = layers.Flatten()(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.Dropout(0.5)(x)\n",
    "        x = layers.Dense(self.latent_dim, activation='relu')(x) # In original: 60\n",
    "\n",
    "        next_shape = self.get_shape_for(stride_factor=4)\n",
    "\n",
    "        # Upsampling. In original: no decoder\n",
    "        x = self.dense_block(x, np.prod(next_shape))\n",
    "        x = layers.Reshape(target_shape=next_shape)(x)\n",
    "        x = self.deconv_block(x, 20, 4, stride=2)\n",
    "        x = self.deconv_block(x, 20, 4, stride=2)\n",
    "        x = self.conv_block(x, 15, 4)\n",
    "        x = self.conv_block(x, 1, 3, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"Pan-2019\")\n",
    "\n",
    "\n",
    "    def build_CNN_2L(self):\n",
    "\n",
    "        # Downsampling\n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        x = self.conv_block(inputs, 32, 3, stride=2, with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, 64, 3, stride=2, with_batchnorm=True, with_dropout=True)\n",
    "        x = layers.Flatten()(x)\n",
    "        x = self.dense_block(x, self.latent_dim, activation='sigmoid')\n",
    "\n",
    "        next_shape = self.get_shape_for(stride_factor=4)\n",
    "\n",
    "        # Upsampling\n",
    "        x = self.dense_block(x, np.prod(next_shape))\n",
    "        x = layers.Reshape(target_shape=next_shape)(x)\n",
    "        x = self.deconv_block(x, 64, 3, stride=2)\n",
    "        x = self.deconv_block(x, 32, 3, stride=2)\n",
    "        x = self.deconv_block(x, 1, 3, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"CNN-v1\")\n",
    "        \n",
    "\n",
    "    def build_convLTSM(self):\n",
    "        \n",
    "        expanded_shape = self.input_size.copy()\n",
    "        expanded_shape.insert(0, 1)\n",
    "        \n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        \n",
    "        x = layers.Reshape(target_shape=expanded_shape)(inputs)\n",
    "        \n",
    "        x = layers.ConvLSTM2D(filters=32, kernel_size=(3, 3), return_sequences=True, padding = 'same',\n",
    "                              go_backwards=True, activation='tanh',\n",
    "                              dropout=0.4, recurrent_dropout=0.2)(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "\n",
    "        x = layers.ConvLSTM2D(filters=16, kernel_size=(3, 3), return_sequences=True, padding = 'same',\n",
    "                              go_backwards=True, activation='tanh',\n",
    "                              dropout=0.4, recurrent_dropout=0.2)(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "\n",
    "        x = layers.ConvLSTM2D(filters=8, kernel_size=(3, 3), return_sequences=False, padding = 'same',\n",
    "                              go_backwards=True, activation='tanh',\n",
    "                              dropout=0.3, recurrent_dropout=0.2)(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "\n",
    "        x = layers.Conv2D(filters=16, kernel_size=(1, 1), activation='relu')(x)\n",
    "\n",
    "        x = layers.MaxPooling2D(pool_size=(2, 2), padding='same')(x)\n",
    "        #model.add(Flatten())\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.Dropout(0.25)(x)\n",
    "\n",
    "        x = layers.Conv2DTranspose(filters=1, kernel_size=(1, 1), strides=(2, 2),  activation='sigmoid',\n",
    "                                   padding='same')(x)\n",
    "\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"Conv-LTSM\")\n",
    "        \n",
    "        \n",
    "    def build_Unet(self):\n",
    "        \"\"\"\n",
    "        Based on: U-Net: https://github.com/nikhilroxtomar/Unet-for-Person-Segmentation/blob/main/model.py\n",
    "        \"\"\"\n",
    "        \n",
    "        # Downsampling\n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        \n",
    "        # Pad if necessary\n",
    "        x = self.padding_block(inputs, factor=2**self.unet_depth)\n",
    "        \n",
    "        skips = []\n",
    "        for i in range(self.unet_depth):\n",
    "            s, x = self.unet_encoder_block(x, self.unet_filters_nb * 2**i)\n",
    "            skips.append(s)\n",
    "        \n",
    "        x = self.conv_block(x, self.unet_filters_nb * 2**self.unet_depth, 3, initializer='he_normal', with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, self.unet_filters_nb * 2**self.unet_depth, 3, initializer='he_normal', with_batchnorm=True)\n",
    "\n",
    "        # Upsampling\n",
    "        for i in reversed(range(self.unet_depth)):\n",
    "            x = self.unet_decoder_block(x, skips[i], self.unet_filters_nb * 2**i, is_last=(i==0))\n",
    "        \n",
    "        # Additional upsampling for downscaling\n",
    "        x = self.handle_output_scaling(x)\n",
    "\n",
    "        x = self.conv_block(x, 1, 1, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"U-Net-v1\")\n",
    "                \n",
    "            \n",
    "    def build_UnetAtt(self):\n",
    "        \"\"\"\n",
    "        U-Net with attention mechanism\n",
    "        \"\"\"\n",
    "        \n",
    "        # Downsampling\n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        \n",
    "        # Pad if necessary\n",
    "        x = self.padding_block(inputs, factor=self.unet_depth**2)\n",
    "        \n",
    "        skips = []\n",
    "        for i in range(self.unet_depth):\n",
    "            s, x = self.unet_encoder_block(x, self.unet_filters_nb * 2**i)\n",
    "            skips.append(s)\n",
    "        \n",
    "        x = self.conv_block(x, self.unet_filters_nb * 2**self.unet_depth, 3, initializer='he_normal', with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, self.unet_filters_nb * 2**self.unet_depth, 3, initializer='he_normal', with_batchnorm=True)\n",
    "\n",
    "        # Upsampling\n",
    "        for i in reversed(range(self.unet_depth)):\n",
    "            x = self.unet_decoder_block_attention(x, skips[i], self.unet_filters_nb * 2**i, is_last=(i==0))\n",
    "        \n",
    "        # Additional upsampling for downscaling\n",
    "        x = self.handle_output_scaling(x)\n",
    "\n",
    "        x = self.conv_block(x, 1, 1, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"U-Net-attention\")\n",
    "        \n",
    "        \n",
    "    def build_RaNet(self):\n",
    "        \"\"\"\n",
    "        Based on: https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2020GL090309\n",
    "        \"\"\"\n",
    "        \n",
    "        expanded_shape = self.input_size.copy()\n",
    "        expanded_shape.insert(len(expanded_shape), int(expanded_shape[-1]/self.levels_nb_3d))\n",
    "        expanded_shape[-2] = self.levels_nb_3d\n",
    "        \n",
    "        print(f\"New shape: {expanded_shape}\")\n",
    "        \n",
    "        inputs = layers.Input(shape=self.input_size)\n",
    "        x = layers.Reshape(target_shape=expanded_shape)(inputs)\n",
    "        \n",
    "        n_filters = 64\n",
    "        \n",
    "        # Downsampling\n",
    "        \n",
    "        #x = self.conv_block(inputs, 32, 3, stride=2, with_batchnorm=True, with_dropout=True)\n",
    "        #x = self.conv_block(x, 64, 3, stride=2, with_batchnorm=True, with_dropout=True)\n",
    "        \n",
    "        x = layers.Conv3D(n_filters, kernel_size=(3, 3, 3),padding=\"same\")(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.MaxPooling3D()(x)\n",
    "        x = layers.Conv3D(n_filters*2, kernel_size=(3, 3, 3),padding=\"same\")(x)\n",
    "        x = layers.BatchNormalization()(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.MaxPooling3D()(x)\n",
    "        x = layers.Conv3D(n_filters*4, kernel_size=(3, 3, 3),padding=\"same\")(x)\n",
    "\n",
    "        x = layers.GlobalAveragePooling3D()(x)\n",
    "        x = layers.Dense(n_filters*4, activation='relu')(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(n_filters*8, activation='relu')(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        x = layers.Dense(n_filters*4, activation='relu')(x)\n",
    "        x = layers.Dropout(0.2)(x)\n",
    "        \n",
    "        next_shape = self.get_shape_for(stride_factor=4)\n",
    "\n",
    "        # Upsampling.\n",
    "        x = self.dense_block(x, np.prod(next_shape))\n",
    "        x = layers.Reshape(target_shape=next_shape)(x)\n",
    "        x = self.deconv_block(x, n_filters*4, 3, stride=2)\n",
    "        x = self.deconv_block(x, n_filters*2, 3, stride=2)\n",
    "        x = self.deconv_block(x, n_filters, 3, stride=2)\n",
    "        x = self.conv_block(x, 1, 3, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "        \n",
    "        # Additional upsampling for downscaling\n",
    "        x = self.handle_output_scaling(x)\n",
    "\n",
    "        x = self.conv_block(x, 1, 1, activation=self.last_activation)\n",
    "        outputs = self.final_cropping_block(x)\n",
    "\n",
    "        self.model = keras.Model(inputs, outputs, name=\"RaNet\")\n",
    "        \n",
    "        \n",
    "    def unet_encoder_block(self, input, filters, kernel_size=3):\n",
    "        x = self.conv_block(input, filters, kernel_size, initializer='he_normal', with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, filters, kernel_size, initializer='he_normal', with_batchnorm=True)\n",
    "        p = layers.MaxPooling2D((2, 2))(x)\n",
    "        \n",
    "        return x, p\n",
    "\n",
    "    \n",
    "    def unet_decoder_block(self, input, skip_features, filters, conv_kernel_size=3, deconv_kernel_size=2, is_last=False):\n",
    "        x = self.deconv_block(input, filters, deconv_kernel_size, stride=2)\n",
    "        x = layers.Concatenate()([x, skip_features])\n",
    "        x = self.conv_block(x, filters, conv_kernel_size, initializer='he_normal', with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, filters, conv_kernel_size, initializer='he_normal', with_batchnorm=(not is_last))\n",
    "\n",
    "        return x\n",
    "        \n",
    "        \n",
    "    def unet_decoder_block_attention(self, input, skip_features, filters, conv_kernel_size=3, deconv_kernel_size=2, is_last=False):\n",
    "        x = attention_up_and_concate(input, skip_features)\n",
    "        x = self.conv_block(x, filters, conv_kernel_size, initializer='he_normal', with_batchnorm=True, with_dropout=True)\n",
    "        x = self.conv_block(x, filters, conv_kernel_size, initializer='he_normal', with_batchnorm=(not is_last))\n",
    "\n",
    "        return x\n",
    "        \n",
    "        \n",
    "    def conv_block(self, input, filters, kernel_size=3, stride=1, padding='same', initializer='default', activation='default', \n",
    "                   with_batchnorm=False, with_pooling=False, with_dropout=False, with_late_activation=False):\n",
    "        if activation == 'default':\n",
    "            activation = self.inner_activation\n",
    "            \n",
    "        conv_activation = activation\n",
    "        if with_late_activation:\n",
    "            conv_activation = None\n",
    "            \n",
    "        if initializer == 'default':\n",
    "            x = layers.Conv2D(filters, kernel_size, strides=(stride, stride), padding=padding, activation=conv_activation)(input)\n",
    "        else:\n",
    "            x = layers.Conv2D(filters, kernel_size, strides=(stride, stride), padding=padding, activation=conv_activation, kernel_initializer=initializer)(input)\n",
    "            \n",
    "        if with_batchnorm:\n",
    "            x = layers.BatchNormalization()(x)\n",
    "        if with_late_activation:\n",
    "            x = layers.Activation(activation)(x)\n",
    "        if with_pooling:\n",
    "            x = layers.MaxPooling2D(pool_size=2)(x)\n",
    "        if with_dropout:\n",
    "            x = layers.SpatialDropout2D(self.dropout_rate)(x)\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    \n",
    "    def deconv_block(self, input, filters, kernel_size=3, stride=1, padding='same', initializer='default', activation='default', \n",
    "                     with_batchnorm=False, with_dropout=False):\n",
    "        if activation == 'default':\n",
    "            activation = self.inner_activation\n",
    "        \n",
    "        if self.use_upsample:\n",
    "            x = layers.UpSampling2D((2, 2))(input)\n",
    "        else:\n",
    "            if initializer == 'default':\n",
    "                x = layers.Conv2DTranspose(filters, kernel_size, strides=stride, padding=padding, activation=activation)(input)\n",
    "            else:\n",
    "                x = layers.Conv2DTranspose(filters, kernel_size, strides=stride, padding=padding, activation=activation, kernel_initializer=initializer)(input)\n",
    "            \n",
    "        if with_batchnorm:\n",
    "            x = layers.BatchNormalization()(x)\n",
    "        if with_dropout:\n",
    "            x = layers.SpatialDropout2D(self.dropout_rate)(x)\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    \n",
    "    def dense_block(self, input, units, activation='default', with_dropout=False):\n",
    "        if activation == 'default':\n",
    "            activation=self.inner_activation\n",
    "            \n",
    "        x = layers.Dense(units, activation=activation)(input)\n",
    "        if with_dropout:\n",
    "            x = layers.Dropout(self.dropout_rate)(x)\n",
    "            \n",
    "        return x\n",
    "\n",
    "    \n",
    "    def handle_output_scaling(self, input, with_batchnorm=False):\n",
    "        if self.output_scaling > 1:\n",
    "            if self.output_scaling == 2:\n",
    "                x = self.deconv_block(input, 64, 3, stride=2, with_batchnorm=with_batchnorm)\n",
    "            elif self.output_scaling == 3:\n",
    "                x = self.deconv_block(input, 64, 3, stride=3, with_batchnorm=with_batchnorm)\n",
    "            elif self.output_scaling == 4:\n",
    "                x = self.deconv_block(input, 64, 3, stride=2, with_batchnorm=with_batchnorm)\n",
    "                x = self.deconv_block(x, 64, 3, stride=2, with_batchnorm=with_batchnorm)\n",
    "            elif self.output_scaling == 5:\n",
    "                x = self.deconv_block(input, 64, 3, stride=3, with_batchnorm=with_batchnorm)\n",
    "                x = self.deconv_block(x, 64, 3, stride=2, with_batchnorm=with_batchnorm)\n",
    "            else:\n",
    "                raise NotImplementedError('Level of downscaling not implemented')\n",
    "        else:\n",
    "            x = input\n",
    "        \n",
    "        if self.output_crop:\n",
    "            raise NotImplementedError('Manual cropping not yet implemented')\n",
    "            \n",
    "        return x\n",
    "            \n",
    "        \n",
    "    def padding_block(self, x, factor):\n",
    "        h, w = x.get_shape().as_list()[1:3]\n",
    "        dh = 0\n",
    "        dw = 0\n",
    "        if h % factor > 0:\n",
    "            dh = factor - h % factor\n",
    "        if w % factor > 0:\n",
    "            dw = factor - w % factor\n",
    "        if dh > 0 or dw > 0:\n",
    "            top_pad = dh//2\n",
    "            bottom_pad = dh//2 + dh%2\n",
    "            left_pad = dw//2\n",
    "            right_pad = dw//2 + dw%2\n",
    "            x = layers.ZeroPadding2D(padding=((top_pad, bottom_pad), (left_pad, right_pad)))(x)\n",
    "        \n",
    "        return x\n",
    "        \n",
    "        \n",
    "    def final_cropping_block(self, x):\n",
    "        # Compute difference between reconstructed width and hight and the desired output size.\n",
    "        h, w = x.get_shape().as_list()[1:3]\n",
    "        h_tgt, w_tgt = self.output_size[:2]\n",
    "        dh = h - h_tgt\n",
    "        dw = w - w_tgt\n",
    "\n",
    "        if dh < 0 or dw < 0:\n",
    "            raise ValueError(f'Negative values in output cropping dh={dh} and dw={dw}')\n",
    "\n",
    "        # Add to decoder cropping layer and final reshaping\n",
    "        x = layers.Cropping2D(cropping=((dh//2, dh-dh//2), (dw//2, dw-dw//2)))(x)\n",
    "        #x = layers.Reshape(target_shape=self.output_size,)(x)\n",
    "        \n",
    "        return x\n",
    "        \n",
    "\n",
    "    def get_shape_for(self, stride_factor):\n",
    "        next_shape = self.output_size.copy()\n",
    "        next_shape[0] = int(np.ceil(next_shape[0]/stride_factor))\n",
    "        next_shape[1] = int(np.ceil(next_shape[1]/stride_factor))\n",
    "\n",
    "        return next_shape\n",
    "\n",
    "        \n",
    "    def call(self, x):\n",
    "        return self.model(x)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cfdc976-5d2f-4c0d-9bb7-bff0b633ef26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computing the necessary output scaling.\n",
    "dlons_x = float(lons_x[1] - lons_x[0])\n",
    "dlats_x = float(lats_x[0] - lats_x[1])\n",
    "dlons_y = float(lons_y[1] - lons_y[0])\n",
    "dlats_y = float(lats_y[0] - lats_y[1])\n",
    "\n",
    "output_scaling = int(dlons_x / dlons_y)\n",
    "output_crop = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edec5d6b-38f3-49f1-9d9e-25b3116d791c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def initiate_optimizer(lr_method, lr=.0004, init_lr=0.01, max_lr=0.01):\n",
    "    if lr_method == 'Cyclical':\n",
    "        # Cyclical learning rate\n",
    "        steps_per_epoch = dg_train.n_samples // BATCH_SIZE\n",
    "        clr = tfa.optimizers.CyclicalLearningRate(\n",
    "            initial_learning_rate=init_lr,\n",
    "            maximal_learning_rate=max_lr,\n",
    "            scale_fn=lambda x: 1/(2.**(x-1)),\n",
    "            step_size=2 * steps_per_epoch)\n",
    "        optimizer = tf.keras.optimizers.Adam(clr)\n",
    "    elif lr_method == 'CosineDecay':\n",
    "        decay_steps = EPOCHS * (dg_train.n_samples / BATCH_SIZE)\n",
    "        lr_decayed_fn = tf.keras.optimizers.schedules.CosineDecay(\n",
    "            init_lr, decay_steps)\n",
    "        optimizer = tf.keras.optimizers.Adam(lr_decayed_fn)\n",
    "    elif lr_method == 'Constant':\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate = lr)\n",
    "    else:\n",
    "        raise ValueError('learning rate schedule not well defined.')\n",
    "        \n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e21809",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute weights for the weighted binary crossentropy\n",
    "weights = class_weight.compute_class_weight(\n",
    "    'balanced',\n",
    "    classes=np.unique(pr_xtrm.values),\n",
    "    y=pr_xtrm.values.flatten()\n",
    ")\n",
    "\n",
    "print('Weights for the weighted binary crossentropy:')\n",
    "print(f'Classes: {np.unique(pr_xtrm.values)}, weights: {weights}')\n",
    "\n",
    "# Create loss function for the extremes\n",
    "xtrm_loss = weighted_binary_cross_entropy(\n",
    "    weights={0: weights[0].astype('float32'), 1: weights[1].astype('float32')})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd0ee543-418c-4371-a4c5-548e71077fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "EPOCHS = 200\n",
    "LR_METHOD = 'Constant'  # Cyclical, CosineDecay, Constant\n",
    "    \n",
    "# Early stopping\n",
    "callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10,\n",
    "                                            restore_best_weights=True)\n",
    "                                            \n",
    "# Default model options\n",
    "opt_model = {'latent_dim': 128,\n",
    "             'dropout_rate': 0.2}\n",
    "\n",
    "# Default training options\n",
    "opt_training = {'epochs': EPOCHS,\n",
    "                'callbacks': [callback]}\n",
    "\n",
    "# Default optimizer options\n",
    "opt_optimizer = {'lr_method': 'Constant',\n",
    "                 'lr': 0.0004,\n",
    "                 'init_lr': 0.01}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a743731-9946-4acf-b3ba-04f8318256c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_report = {\n",
    "          'Dav-orig': {'model': 'Davenport-2021', 'run': False,\n",
    "                       'opt_model': {'latent_dim': 16},\n",
    "                       'opt_optimizer': {'lr_method': 'Constant', 'lr': 4e-4}}, # original\n",
    "          'Dav-64': {'model': 'Davenport-2021', 'run': False,\n",
    "                     'opt_model': {'latent_dim': 64},\n",
    "                     'opt_optimizer': {'lr_method': 'Constant', 'lr': 4e-4}},\n",
    "          'Pan-orig': {'model': 'Pan-2019', 'run': False,\n",
    "                       'opt_model': {'latent_dim': 60},\n",
    "                       'opt_optimizer': {'lr_method': 'Constant', 'lr': 1e-4}},\n",
    "          'CNN-2l': {'model': 'CNN-2L', 'run': False,\n",
    "                     'opt_model': {'latent_dim': 64},\n",
    "                     'opt_optimizer': {'lr_method': 'Constant', 'lr': 1e-3}}, # higher lr\n",
    "          'UNET2': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 2, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6605599-ba82-4404-899a-684d591be127",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_unets = {\n",
    "          'UNET1': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 1, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "          'UNET2': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 2, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "          'UNET3': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 3, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "          'UNET4': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 4, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "          'UNET-att': {'model': 'Unet-attention', 'run': False,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eda3b0d2-ca68-4b90-aa8a-e2d592413e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_unet2 = {\n",
    "          'UNET2': {'model': 'Unet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 2, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93e8016-9b5f-4224-95df-2f21606353c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_ranet = {\n",
    "          'RaNet': {'model': 'RaNet', 'run': True,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d02e517d-11b1-4edc-ac8f-0cef1a17bc53",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = model_unet2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "231ef4a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict the amount of precipitation\n",
    "df_prec = pd.DataFrame(columns = ['id', 'name', 'n_params', 'opt_model', 'opt_optimizer',\n",
    "                                  'train_pr_rmse', 'test_pr_rmse', \n",
    "                                  'train_xtrm_precision', 'test_xtrm_precision', \n",
    "                                  'train_xtrm_recall', 'test_xtrm_recall'])\n",
    "df_xtrm = pd.DataFrame(columns = ['id', 'name', 'n_params', 'opt_model', 'opt_optimizer',\n",
    "                                  'train_xtrm_roc_auc', 'test_xtrm_roc_auc', \n",
    "                                  'train_xtrm_precision', 'test_xtrm_precision', \n",
    "                                  'train_xtrm_recall', 'test_xtrm_recall'])\n",
    "\n",
    "train_for_prec = True\n",
    "train_for_xtrm = True\n",
    "history_log_level = 1\n",
    "\n",
    "if CREATE_MASK_EOBS:\n",
    "    loss_regression = 'mse_nans'\n",
    "else:\n",
    "    loss_regression = 'mse'\n",
    "    \n",
    "\n",
    "models_prec = []\n",
    "models_xtrm = []\n",
    "\n",
    "if train_for_prec:\n",
    "        \n",
    "    for m_id in models:\n",
    "        # Clear session and set tf seed\n",
    "        keras.backend.clear_session()\n",
    "        tf.random.set_seed(42)\n",
    "        \n",
    "        if not models[m_id]['run']:\n",
    "            continue\n",
    "\n",
    "        # Extract model name and options\n",
    "        model = models[m_id]['model']\n",
    "        opt_model_i = models[m_id]['opt_model']\n",
    "        opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "        opt_model_new = opt_model.copy()\n",
    "        opt_model_new.update(opt_model_i)\n",
    "        opt_optimizer_new = opt_optimizer.copy()\n",
    "        opt_optimizer_new.update(opt_optimizer_i)\n",
    "        print(f'Running: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "        df_prec = df_prec.append({'id': m_id, 'name': model, 'opt_model': opt_model_i, 'opt_optimizer': opt_optimizer_i}, ignore_index=True)\n",
    "\n",
    "        # Switch to precipitation values\n",
    "        dg_train.for_extremes(False)\n",
    "        dg_valid.for_extremes(False)\n",
    "        dg_test.for_extremes(False)\n",
    "        \n",
    "        optimizer = initiate_optimizer(**opt_optimizer_new)\n",
    "\n",
    "        # Load if previously saved\n",
    "        tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "            pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "        hashed_name = f'prec_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "        tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "        \n",
    "        if tmp_file.exists():\n",
    "            m = keras.models.load_model(tmp_file, custom_objects = {\"MeanSquaredErrorNans\": MeanSquaredErrorNans})\n",
    "            \n",
    "        else:\n",
    "            # Create the model and compile\n",
    "            m = DeepFactory(model, i_shape, o_shape, for_extremes=False, **opt_model_new)\n",
    "            # Warning: When using regularizers, the loss function is the entire loss, ie (loss metrics) + (regularization term)!\n",
    "            # But the loss displayed as part of the metrics, is only the loss metric. The regularization term is not added there. -> can be different!!\n",
    "            loss_fct = 'mse'\n",
    "            if loss_regression == 'mse_nans':\n",
    "                loss_fct = MeanSquaredErrorNans()\n",
    "            m.compile(\n",
    "                loss=loss_fct, \n",
    "                metrics=[loss_fct], \n",
    "                optimizer=optimizer\n",
    "            )\n",
    "            print(f'Number of parameters: {m.model.count_params()}')\n",
    "\n",
    "            # Train\n",
    "            hist = m.fit(dg_train, validation_data=dg_valid, verbose=history_log_level, **opt_training)\n",
    "            \n",
    "            # Plot training evolution\n",
    "            pd.DataFrame(hist.history).plot(figsize=(8, 5))\n",
    "            ymin = 0.8 * min([min(hist.history['loss']), min(hist.history['val_loss']), min(hist.history[loss_regression])])\n",
    "            ymax = min(max(max(hist.history['loss']), max(hist.history['val_loss'])), \n",
    "                       2 * max(sum(hist.history['loss'])/len(hist.history['loss']), sum(hist.history['val_loss'])/len(hist.history['val_loss'])))\n",
    "            plt.ylim(ymin, ymax)\n",
    "            plt.grid(True)\n",
    "            plt.title(model)\n",
    "            plt.tight_layout()\n",
    "            now = datetime.datetime.now()\n",
    "            plt.savefig(f'figures/training_prec_{m_id}_{now.strftime(\"%Y-%m-%d_%H-%M-%S\")}.pdf')\n",
    "            plt.savefig(f'figures/training_prec_{m_id}_{now.strftime(\"%Y-%m-%d_%H-%M-%S\")}.png')\n",
    "            plt.show()\n",
    "\n",
    "            # Saving the model\n",
    "            m.save(f'tmp/{hashed_name}')\n",
    "        \n",
    "        df_prec.at[df_prec.index[-1], 'n_params'] = m.model.count_params()\n",
    "        models_prec.append(m)\n",
    "        \n",
    "        # Predict and save scores\n",
    "        y_pred_train = m.predict(dg_train.X.to_numpy()).squeeze()\n",
    "        y_pred_test = m.predict(dg_test.X.to_numpy()).squeeze()\n",
    "        df_prec.at[df_prec.index[-1], 'train_pr_rmse'] = np.sqrt(np.nanmean(np.square(np.subtract(dg_train.y.to_numpy().squeeze(), y_pred_train))))\n",
    "        df_prec.at[df_prec.index[-1], 'test_pr_rmse'] =  np.sqrt(np.nanmean(np.square(np.subtract(dg_test.y.to_numpy().squeeze(), y_pred_test))))\n",
    "        \n",
    "        # save predictions\n",
    "        np.save(f'tmp/data/predictions/y_pred_train_{m_id}_{int(PRECIP_XTRM*100)}.npy', y_pred_train)\n",
    "        np.save(f'tmp/data/predictions/y_pred_test_{m_id}_{int(PRECIP_XTRM*100)}.npy', y_pred_test)\n",
    "        \n",
    "        # Analyze predictions\n",
    "        print('Plotting results of the training period, amount to precipitation.')\n",
    "        precision, recall = analyze_predictions(y_pred_train, dg_train, qq, lons_y, lats_y, mask=mask, pred_xtrm=False, show_plots=False)\n",
    "        df_prec.at[df_prec.index[-1], 'train_xtrm_precision'] = precision\n",
    "        df_prec.at[df_prec.index[-1], 'train_xtrm_recall'] = recall\n",
    "        plt.show()\n",
    "\n",
    "        print('Plotting results of the testing period.')\n",
    "        precision, recall = analyze_predictions(y_pred_test, dg_test, qq, lons_y, lats_y, mask=mask, pred_xtrm=False)\n",
    "        df_prec.at[df_prec.index[-1], 'test_xtrm_precision'] = precision\n",
    "        df_prec.at[df_prec.index[-1], 'test_xtrm_recall'] = recall\n",
    "        plt.show()\n",
    "    \n",
    "    \n",
    "        print(df_prec.iloc[-1])\n",
    "        \n",
    "        print(f\"\\n{'*' * 100}\\n\")\n",
    "        \n",
    "if train_for_xtrm:\n",
    "\n",
    "    for m_id in models:\n",
    "        # Clear session and set tf seed\n",
    "        keras.backend.clear_session()\n",
    "        tf.random.set_seed(42)\n",
    "\n",
    "        if not models[m_id]['run']:\n",
    "            continue\n",
    "        \n",
    "        # Extract model name and options\n",
    "        model = models[m_id]['model']\n",
    "        opt_model_i = models[m_id]['opt_model']\n",
    "        opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "        opt_model_new = opt_model.copy()\n",
    "        opt_model_new.update(opt_model_i)\n",
    "        opt_optimizer_new = opt_optimizer.copy()\n",
    "        opt_optimizer_new.update(opt_optimizer_i)\n",
    "        print(f'Running: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "        df_xtrm = df_xtrm.append({'id': m_id, 'name': model, 'opt_model': opt_model_i, 'opt_optimizer': opt_optimizer_i}, ignore_index=True)\n",
    "        \n",
    "        # Switch to precipitation extremes\n",
    "        dg_train.for_extremes(True)\n",
    "        dg_valid.for_extremes(True)\n",
    "        dg_test.for_extremes(True)\n",
    "        \n",
    "        optimizer = initiate_optimizer(**opt_optimizer_new)\n",
    "        \n",
    "        # Load if previously saved\n",
    "        tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "            pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "        hashed_name = f'xtrm_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "        tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "        \n",
    "        if tmp_file.exists():\n",
    "            custom_objects = {\"weighted_cross_entropy_fn\": weighted_binary_cross_entropy}\n",
    "            with keras.utils.custom_object_scope(custom_objects):\n",
    "                m = keras.models.load_model(tmp_file)\n",
    "            \n",
    "        else:\n",
    "            # Create the model and compile\n",
    "            m = DeepFactory(model, i_shape, o_shape, for_extremes=True, **opt_model_new)\n",
    "            m.compile(\n",
    "                loss=xtrm_loss,\n",
    "                optimizer=optimizer\n",
    "            )\n",
    "            print(f'Number of parameters: {m.model.count_params()}')\n",
    "\n",
    "            # Train\n",
    "            hist = m.fit(dg_train, validation_data=dg_valid, verbose=history_log_level, **opt_training)\n",
    "            \n",
    "            # Plot training evolution\n",
    "            pd.DataFrame(hist.history).plot(figsize=(8, 5))\n",
    "            ymin = 0.8 * min([min(hist.history['loss']), min(hist.history['val_loss'])])\n",
    "            ymax = min(max(max(hist.history['loss']), max(hist.history['val_loss'])), \n",
    "                       2 * max(sum(hist.history['loss'])/len(hist.history['loss']), sum(hist.history['val_loss'])/len(hist.history['val_loss'])))\n",
    "            plt.ylim(ymin, ymax)\n",
    "            plt.grid(True)\n",
    "            plt.title(model)\n",
    "            plt.tight_layout()\n",
    "            now = datetime.datetime.now()\n",
    "            plt.savefig(f'figures/training_xtrm_{m_id}_{now.strftime(\"%Y-%m-%d_%H-%M-%S\")}.pdf')\n",
    "            plt.savefig(f'figures/training_xtrm_{m_id}_{now.strftime(\"%Y-%m-%d_%H-%M-%S\")}.png')\n",
    "            plt.show()\n",
    "\n",
    "            # Saving the model\n",
    "            m.save(f'tmp/{hashed_name}')\n",
    "        \n",
    "        df_xtrm.at[df_xtrm.index[-1], 'n_params'] = m.model.count_params()\n",
    "        models_xtrm.append(m)\n",
    "        \n",
    "        # Assess and save scores\n",
    "        y_pred_train = m.predict(dg_train.X.to_numpy()).squeeze()\n",
    "        y_pred_test = m.predict(dg_test.X.to_numpy()).squeeze()\n",
    "        \n",
    "        # save predictions\n",
    "        np.save(f'tmp/data/predictions/y_pred_train_xtrm_{m_id}_{int(PRECIP_XTRM*100)}.npy', y_pred_train)\n",
    "        np.save(f'tmp/data/predictions/y_pred_test_xtrm_{m_id}_{int(PRECIP_XTRM*100)}.npy', y_pred_test)\n",
    "        \n",
    "        #Analyze predictions\n",
    "        print('Plotting results of the training period.')\n",
    "        precision, recall, roc_auc = analyze_predictions(y_pred_train, dg_train, qq, lons_y, lats_y, mask=mask, pred_xtrm=True, show_plots=False)\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_precision'] = precision\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_recall'] = recall\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_roc_auc'] = roc_auc\n",
    "        plt.show()\n",
    "        \n",
    "        print('Plotting results of the testing period.')\n",
    "        precision, recall, roc_auc = analyze_predictions(y_pred_test, dg_test, qq, lons_y, lats_y, mask=mask, pred_xtrm=True)\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'test_xtrm_precision'] = precision\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'test_xtrm_recall'] = recall\n",
    "        df_xtrm.at[df_xtrm.index[-1], 'test_xtrm_roc_auc'] = roc_auc\n",
    "        plt.show()\n",
    "\n",
    "        print(df_xtrm.iloc[-1])\n",
    "        \n",
    "        print(f\"\\n{'*' * 100}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a495bdc0-fc28-4433-9b28-4e9627a67d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_prec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5df59f7-5926-46ac-be5d-f7ca5b71d493",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_xtrm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c037df7-37e2-4ceb-a83a-373a6bde453e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting\n",
    "n_models = 0\n",
    "for m_id in models:\n",
    "    if models[m_id]['run']:\n",
    "        n_models += 1\n",
    "        \n",
    "n_rows = n_models+1\n",
    "fig, axs = plt.subplots(n_rows, 3, figsize=(10, n_rows*3))\n",
    "\n",
    "# Extract true values\n",
    "y_xtrm = dg_test.y_xtrm.to_numpy().squeeze()\n",
    "y_prec = dg_test.y.to_numpy().squeeze()\n",
    "\n",
    "# Get the index of the max # of extremes\n",
    "i_max_obs = np.argmax(np.sum(y_xtrm, axis=(1,2)))\n",
    "\n",
    "vmax = np.max(y_prec[i_max_obs])\n",
    "\n",
    "plot_map(axs[0, 0], lons_y, lats_y, y_prec[i_max_obs], title=\"Prec. value - truth\", vmin=0, vmax=vmax, show_colorbar=False, cmap=mpl.cm.YlGnBu)\n",
    "plot_map(axs[0, 1], lons_y, lats_y, y_xtrm[i_max_obs], title=\"Prec. extreme - truth\", vmin=0, vmax=1, show_colorbar=False)\n",
    "plot_map(axs[0, 2], lons_y, lats_y, y_xtrm[i_max_obs], title=\"Prec. extreme - truth\", vmin=0, vmax=1, show_colorbar=False)\n",
    "    \n",
    "for idx, m_id in enumerate(models):\n",
    "\n",
    "    if not models[m_id]['run']:\n",
    "        continue\n",
    "\n",
    "    # Precip value - extract model name and options\n",
    "    model = models[m_id]['model']\n",
    "    opt_model_i = models[m_id]['opt_model']\n",
    "    opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "    opt_model_new = opt_model.copy()\n",
    "    opt_model_new.update(opt_model_i)\n",
    "    opt_optimizer_new = opt_optimizer.copy()\n",
    "    opt_optimizer_new.update(opt_optimizer_i)\n",
    "    print(f'Plotting: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "\n",
    "    # Switch to precipitation values\n",
    "    dg_test.for_extremes(False)\n",
    "\n",
    "    # Load if previously saved\n",
    "    tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "        pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "    hashed_name = f'prec_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "    tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "\n",
    "    m = keras.models.load_model(tmp_file, custom_objects = {\"MeanSquaredErrorNans\": MeanSquaredErrorNans})\n",
    "\n",
    "    # Predict and save scores\n",
    "    y_pred = m.predict(dg_test.X[i_max_obs:i_max_obs+2].to_numpy()).squeeze()\n",
    "    \n",
    "    y_pred_bool = y_pred > qq.to_numpy().squeeze()\n",
    "\n",
    "    # Multiply to transorm to numeric values\n",
    "    y_pred_bool = y_pred_bool * 1\n",
    "    \n",
    "    plot_map(axs[idx+1, 0], lons_y, lats_y, y_pred[0], title=f\"Prec. value - {m_id}\", vmin=0, vmax=vmax, show_colorbar=False, cmap=mpl.cm.YlGnBu)\n",
    "    plot_map(axs[idx+1, 1], lons_y, lats_y, y_pred_bool[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "\n",
    "    # Preci extremes - extract model name and options\n",
    "    model = models[m_id]['model']\n",
    "    opt_model_i = models[m_id]['opt_model']\n",
    "    opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "    opt_model_new = opt_model.copy()\n",
    "    opt_model_new.update(opt_model_i)\n",
    "    opt_optimizer_new = opt_optimizer.copy()\n",
    "    opt_optimizer_new.update(opt_optimizer_i)\n",
    "    print(f'Plotting: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "\n",
    "    # Switch to precipitation extremes\n",
    "    dg_test.for_extremes(True)\n",
    "\n",
    "    # Load if previously saved\n",
    "    tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "        pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "    hashed_name = f'xtrm_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "    tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "\n",
    "    custom_objects = {\"weighted_cross_entropy_fn\": weighted_binary_cross_entropy, \"MeanSquaredErrorNans\": MeanSquaredErrorNans}\n",
    "    with keras.utils.custom_object_scope(custom_objects):\n",
    "        m = keras.models.load_model(tmp_file)\n",
    "\n",
    "    # Assess and save scores\n",
    "    y_pred_xtrm = m.predict(dg_test.X[i_max_obs:i_max_obs+2].to_numpy()).squeeze()\n",
    "\n",
    "    plot_map(axs[idx+1, 2], lons_y, lats_y, y_pred_xtrm[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "#plt.savefig('figures/plot_model_comparison.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451288bf-8449-4d11-a2c3-de7fb90caa90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the baseline\n",
    "n_rows = n_models\n",
    "fig, axs = plt.subplots(n_rows, 3, figsize=(10, n_rows*3))\n",
    "\n",
    "baseline=0  # 0 or 0.5\n",
    "\n",
    "if baseline==0:\n",
    "    baseline_inputs = np.zeros(dg_test.X[i_max_obs:i_max_obs+2].to_numpy().shape)\n",
    "elif baseline==0.5:\n",
    "    baseline_inputs = np.ones(dg_test.X[i_max_obs:i_max_obs+2].to_numpy().shape) * 0.5\n",
    "elif baseline==1:\n",
    "    baseline_inputs = np.ones(dg_test.X[i_max_obs:i_max_obs+2].to_numpy().shape)\n",
    "\n",
    "    \n",
    "for idx, m_id in enumerate(models):\n",
    "\n",
    "    if not models[m_id]['run']:\n",
    "        continue\n",
    "\n",
    "    # Precip value - extract model name and options\n",
    "    model = models[m_id]['model']\n",
    "    opt_model_i = models[m_id]['opt_model']\n",
    "    opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "    opt_model_new = opt_model.copy()\n",
    "    opt_model_new.update(opt_model_i)\n",
    "    opt_optimizer_new = opt_optimizer.copy()\n",
    "    opt_optimizer_new.update(opt_optimizer_i)\n",
    "    print(f'Plotting: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "\n",
    "    # Switch to precipitation values\n",
    "    dg_test.for_extremes(False)\n",
    "\n",
    "    # Load if previously saved\n",
    "    tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "        pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "    hashed_name = f'prec_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "    tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "\n",
    "    m = keras.models.load_model(tmp_file, custom_objects = {\"MeanSquaredErrorNans\": MeanSquaredErrorNans})\n",
    "\n",
    "    # Predict\n",
    "    y_pred = m.predict(baseline_inputs).squeeze()\n",
    "    y_pred_bool = y_pred > qq.to_numpy().squeeze()\n",
    "    \n",
    "    print(f\"Mean precipitation: {y_pred.mean()}\")\n",
    "\n",
    "    # Multiply to transorm to numeric values\n",
    "    y_pred_bool = y_pred_bool * 1\n",
    "    \n",
    "    if n_rows==1:\n",
    "        plot_map(axs[0], lons_y, lats_y, y_pred[0], title=f\"Prec. value - {m_id} (max: {y_pred.max():.1f}mm)\", vmin=0, vmax=y_pred.max(), show_colorbar=False, cmap=mpl.cm.YlGnBu)\n",
    "        plot_map(axs[1], lons_y, lats_y, y_pred_bool[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "    else:\n",
    "        plot_map(axs[idx, 0], lons_y, lats_y, y_pred[0], title=f\"Prec. value - {m_id}\", vmin=0, vmax=vmax, show_colorbar=False, cmap=mpl.cm.YlGnBu)\n",
    "        plot_map(axs[idx, 1], lons_y, lats_y, y_pred_bool[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "\n",
    "    # Preci extremes - extract model name and options\n",
    "    model = models[m_id]['model']\n",
    "    opt_model_i = models[m_id]['opt_model']\n",
    "    opt_optimizer_i = models[m_id]['opt_optimizer']\n",
    "    opt_model_new = opt_model.copy()\n",
    "    opt_model_new.update(opt_model_i)\n",
    "    opt_optimizer_new = opt_optimizer.copy()\n",
    "    opt_optimizer_new.update(opt_optimizer_i)\n",
    "    print(f'Plotting: {m_id} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "\n",
    "    # Switch to precipitation extremes\n",
    "    dg_test.for_extremes(True)\n",
    "\n",
    "    # Load if previously saved\n",
    "    tag = pickle.dumps(opt_model_new) + pickle.dumps(opt_optimizer_new) + \\\n",
    "        pickle.dumps(i_shape + o_shape) + pickle.dumps(PRECIP_XTRM)\n",
    "    hashed_name = f'xtrm_{m_id}_{hashlib.md5(tag).hexdigest()}'\n",
    "    tmp_file = pathlib.Path(f'tmp/{hashed_name}')\n",
    "\n",
    "    custom_objects = {\"weighted_cross_entropy_fn\": weighted_binary_cross_entropy, \"MeanSquaredErrorNans\": MeanSquaredErrorNans}\n",
    "    with keras.utils.custom_object_scope(custom_objects):\n",
    "        m = keras.models.load_model(tmp_file)\n",
    "\n",
    "    # Predict\n",
    "    y_pred_xtrm = m.predict(baseline_inputs).squeeze()\n",
    "\n",
    "    if n_rows==1:\n",
    "        plot_map(axs[2], lons_y, lats_y, y_pred_xtrm[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "    else:\n",
    "        plot_map(axs[idx, 2], lons_y, lats_y, y_pred_xtrm[0], title=f\"Prec. extreme - {m_id}\", vmin=0, vmax=1, show_colorbar=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(f'figures/plot_model_baseline_{baseline}.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f97de9d-a95e-45d8-b47f-14c17ad5bcd3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "63fd5069d213b44bf678585dea6b12cceca9941eaf7f819626cde1f2670de90d"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}