UNet_inputsubsets_comparison.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0b2bbd3-440e-4c52-90fd-51fc4e484f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This notebook assesses the number of inputs to improve the performance of the best UNET\n",
    "# The inputs are sub-sequently added 5-by-5 and the metrics are then calculated\n",
    "\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",
    "# 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",
    "# 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.DNN_models import *\n",
    "\n",
    "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\n",
    "CREATE_MASK_EOBS = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "764a6a7b-ac45-45d0-8fb2-1d3256e4c4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "conf = yaml.safe_load(open(\"config.yaml\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54a16f9e-0603-4e49-9f88-1430b1f79b35",
   "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",
    "    #pr = pr.fillna(0) # Over the seas. Not optimal...\n",
    "else:\n",
    "    raise ValueError('Precipitation data not well defined')\n",
    "\n",
    "    # Add a dimension to be used as channel in the DNN\n",
    "pr = pr.expand_dims('level', -1)\n",
    "\n",
    "\n",
    "# Compute the extreme exceedence\n",
    "qq = xr.DataArray(pr).quantile(PRECIP_XTRM, dim='time')\n",
    "pr_xtrm = xr.DataArray(pr > qq)\n",
    "pr_xtrm = pr_xtrm*1 # Transform to number\n",
    "\n",
    "# Extract coordinates for precip\n",
    "lats_y = pr.lat.to_numpy()\n",
    "lons_y = pr.lon.to_numpy()\n",
    "\n",
    "# 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()\n",
    "    \n",
    "    \n",
    "\n",
    "## Input data: meteorological fields\n",
    "# Load geopotential height\n",
    "print('loading predictors')\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ea2cfc-3db9-4b16-8aeb-9c1c43d48962",
   "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",
    "\n",
    "# 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": "code",
   "execution_count": null,
   "id": "b9447416-a0b0-4eee-b721-37d511de48c8",
   "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'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5a8280-b6dd-4a74-b25f-3a71f5273e6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#rain the model for each subset\n",
    "subsets = ['subset_1','subset_2','subset_3','subset_4','subset_5','subset_6']\n",
    "YY_VALID = 2005"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b98375-9a70-4ede-acc5-7bd566fa720e",
   "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": "2a2a05c2-8867-44da-85f9-3b07e871f5c0",
   "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": "6ad9c0e2-0e33-45e0-89c0-461cd064304e",
   "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}\n",
    "loss_regression = 'mse'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d9b496-83bc-4596-a84f-0bccb72992d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_unets = {\n",
    "          'UNET1': {'model': 'Unet', 'run': False,\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': False,\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': False,\n",
    "                   'opt_model': {'output_scaling': output_scaling, 'output_crop': output_crop, 'unet_depth': 4, 'use_upsample': False},\n",
    "                   'opt_optimizer': {'lr_method': 'Constant'}},\n",
    "         }\n",
    "\n",
    "\n",
    "models = models_unets\n",
    "\n",
    "train_for_prec = True\n",
    "train_for_xtrm = True\n",
    "history_log_level = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54d8a71a-564f-4430-92ab-39019fb7d7be",
   "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'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18509255-d5b4-433d-a5b0-493af3d47e7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_prec = []\n",
    "models_xtrm = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9b0834b-936d-4418-b18c-7fbd99067122",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sub in subsets:\n",
    "    \n",
    "    print(conf[sub])\n",
    "    dic = conf[sub]\n",
    "    print(dic)\n",
    "    \n",
    "    if 'tcwv' in dic.keys(): # make sure to convert to nontype\n",
    "        dic['tcwv']=None\n",
    "    # apply datagenerator\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",
    "\n",
    "\n",
    "    i_shape = dg_train.X.shape[1:]\n",
    "    o_shape = dg_train.y.shape[1:]\n",
    "\n",
    "    print(f'X shape: {i_shape}')\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} - {sub} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "            df_prec = df_prec.append({'id': sub, '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",
    "            m = DeepFactory_Keras(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.model.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.model.fit(dg_train, validation_data=dg_valid, verbose=history_log_level, **opt_training)\n",
    "\n",
    "             \n",
    "                # Saving the model\n",
    "            m.model.save_weights(f'tmp/keras/{PRECIP_DATA}_{PRECIP_XTRM}_{m_id}_{sub}.h5')\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.model.predict(dg_train.X.to_numpy()).squeeze()\n",
    "            y_pred_test = m.model.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",
    "            # Analyze predictions\n",
    "            print('Plotting results of the training period, amount to precipitation.')\n",
    "            precision, recall = analyze_predictions(y_pred_train, dg_train, qq, 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, mask=mask, pred_xtrm=False,  show_plots=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",
    "            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} - {sub} - {model} - {opt_model_i} - {opt_optimizer_i}')\n",
    "                df_xtrm = df_xtrm.append({'id': sub, '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",
    "                m = DeepFactory_Keras(model, i_shape, o_shape, for_extremes=True, **opt_model_new)\n",
    "                m.model.compile(\n",
    "                        loss=xtrm_loss,\n",
    "                        optimizer=optimizer)\n",
    "                print(f'Number of parameters: {m.model.count_params()}')\n",
    "\n",
    "                    # Train\n",
    "                hist = m.model.fit(dg_train, validation_data=dg_valid, verbose=history_log_level, **opt_training)\n",
    "\n",
    "\n",
    "                # Saving the model\n",
    "                m.model.save_weights(f'tmp/keras/{PRECIP_DATA}_{PRECIP_XTRM}_{m_id}_{sub}_xtrm.h5')\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.model.predict(dg_train.X.to_numpy()).squeeze()\n",
    "                y_pred_test = m.model.predict(dg_test.X.to_numpy()).squeeze()\n",
    "\n",
    "                # save predictions\n",
    "                #np.save(f'tmp/data/predictions/y_pred_train_xtrm_{m_id}.npy',y_pred_train)\n",
    "                #np.save(f'tmp/data/predictions/y_pred_test_xtrm_{m_id}.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, 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, mask=mask, pred_xtrm=True, show_plots=False)\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": "bf16813a-d996-40a1-8372-804d0c104331",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_xtrm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a84495-93b6-47f4-bd86-3bbb7e90fb58",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_prec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05cb453f-cedd-405f-ae86-c1bac85c140d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_prec.to_csv('tmp/df_prec_Unet2_inputtest.csv')\n",
    "df_xtrm.to_csv('tmp/df_xtrm_Unet2_inputtest.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a733a91f-e315-4da4-86bb-a179aa50696d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}