{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4d872421-be6b-43cd-ac61-158ef7170c0f",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3b75a90-4ff0-4b3e-ad02-b3e838c502aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# builtins\n",
    "import datetime\n",
    "import warnings\n",
    "import calendar\n",
    "\n",
    "# externals\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import scipy.stats as stats\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "import scipy.stats as stats\n",
    "from IPython.display import Image\n",
    "from sklearn.metrics import r2_score, roc_curve, auc, classification_report\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# locals\n",
    "from climax.main.io import ERA5_PATH, OBS_PATH, TARGET_PATH, DEM_PATH\n",
    "from climax.main.config import CALIB_PERIOD, VALID_PERIOD\n",
    "from pysegcnn.core.utils import search_files\n",
    "from pysegcnn.core.graphics import plot_classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741ccf7a-80e1-46bb-926e-a68c14f0c9b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# entire reference period\n",
    "REFERENCE_PERIOD = np.concatenate([CALIB_PERIOD, VALID_PERIOD], axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efaaf766-e1d3-48a7-8702-7fecde6c8ca3",
   "metadata": {},
   "source": [
    "### Load observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8049468-d848-4dd3-9082-97ad51e5f5ff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# model predictions and observations NetCDF \n",
    "y_true_pr = xr.open_dataset(search_files(OBS_PATH.joinpath('pr'), 'OBS_pr(.*).nc$').pop())\n",
    "y_true_tmax = xr.open_dataset(search_files(OBS_PATH.joinpath('tasmax'), 'OBS_tasmax(.*).nc$').pop())\n",
    "y_true_tmin = xr.open_dataset(search_files(OBS_PATH.joinpath('tasmin'), 'OBS_tasmin(.*).nc$').pop())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69390143-4c30-443f-9ef5-d1dcde4b2592",
   "metadata": {},
   "source": [
    "### Load ERA-5 reference dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1481f20-2d91-429c-b38d-a6f74f97072d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# search ERA-5 reference dataset\n",
    "y_refe_pr = xr.open_dataset(search_files(ERA5_PATH.joinpath('ERA5', 'total_precipitation'), '.nc$').pop())\n",
    "y_refe_tmax = xr.open_dataset(search_files(ERA5_PATH.joinpath('ERA5', '2m_{}_temperature'.format('max')), '.nc$').pop())\n",
    "y_refe_tmin = xr.open_dataset(search_files(ERA5_PATH.joinpath('ERA5', '2m_{}_temperature'.format('min')), '.nc$').pop())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45e11517-197d-4b9c-bd9f-9f96a11ded98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert to °C\n",
    "y_refe_tmax = y_refe_tmax - 273.15\n",
    "y_refe_tmin = y_refe_tmin - 273.15"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04315e91-2099-4ca7-a324-7173bdcf7750",
   "metadata": {},
   "source": [
    "### Select time period"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "437711c9-94db-4f02-82d1-8dabd109931b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# time period\n",
    "PERIOD = REFERENCE_PERIOD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b7df28-7490-4d2f-99e3-660f1dc6e99f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset observations to time period\n",
    "y_true_pr = y_true_pr.sel(time=PERIOD).precipitation\n",
    "y_true_tmax = y_true_tmax.sel(time=PERIOD).tasmax\n",
    "y_true_tmin = y_true_tmin.sel(time=PERIOD).tasmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b5b7d6-63ea-4cb9-955f-4df9e0b53789",
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset Era-5 to time period\n",
    "y_refe_pr = y_refe_pr.sel(time=PERIOD).drop_vars('lambert_azimuthal_equal_area').rename({'tp': 'precipitation'}).precipitation\n",
    "y_refe_tmax = y_refe_tmax.sel(time=PERIOD).drop_vars('lambert_azimuthal_equal_area').rename({'t2m': 'tasmax'}).tasmax\n",
    "y_refe_tmin = y_refe_tmin.sel(time=PERIOD).drop_vars('lambert_azimuthal_equal_area').rename({'t2m': 'tasmin'}).tasmin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98d1b503-1f1a-4a8e-a8e4-eac1e045f320",
   "metadata": {},
   "source": [
    "## Align datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef839050-2355-4da2-ad19-1e43c4b46414",
   "metadata": {},
   "outputs": [],
   "source": [
    "# precipitation\n",
    "y_true_pr, y_refe_pr = xr.align(y_true_pr, y_refe_pr, join='override')\n",
    "y_refe_pr = y_refe_pr.where(~np.isnan(y_true_pr), other=np.nan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f628cf0c-5465-48e2-9d83-89ec407783e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tasmax\n",
    "y_true_tmax, y_refe_tmax = xr.align(y_true_tmax, y_refe_tmax, join='override')\n",
    "y_refe_tmax = y_refe_tmax.where(~np.isnan(y_true_tmax), other=np.nan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eaa502f-5865-402f-810b-664c1275999f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tasmin\n",
    "y_true_tmin, y_refe_tmin = xr.align(y_true_tmin, y_refe_tmin, join='override')\n",
    "y_refe_tmin = y_refe_tmin.where(~np.isnan(y_true_tmin), other=np.nan)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9045be71-d0ef-4457-9655-68b1f85d5cfe",
   "metadata": {},
   "source": [
    "### Plot ERA-5 vs. Observed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58fa1934-c9fc-4e42-9f2b-942a5ccf5617",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_refe_values = y_refe_pr.resample(time='1M').sum(skipna=False).groupby('time.month').mean(dim='time')\n",
    "y_true_values = y_true_pr.resample(time='1M').sum(skipna=False).groupby('time.month').mean(dim='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40be10b7-c868-4cce-b1fb-d201366c17c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "bias_pr = ((y_refe_values - y_true_values) / y_true_values) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c52eadeb-1e46-41e6-a7b7-9178f3d65536",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot average of observation and reference\n",
    "vmin, vmax = 0, 150\n",
    "fig, axes = plt.subplots(1, 3, figsize=(24, 8), sharex=True, sharey=True)\n",
    "axes = axes.flatten()\n",
    "\n",
    "# plot Era-5 reanalysis\n",
    "im1 = axes[0].imshow(y_refe_values.mean(dim='month'), origin='lower', cmap='viridis_r', vmin=vmin, vmax=vmax)\n",
    "im2 = axes[1].imshow(y_true_values.mean(dim='month'), origin='lower', cmap='viridis_r', vmin=vmin, vmax=vmax)\n",
    "im3 = axes[2].imshow(bias_pr.mean(dim='month'), origin='lower', cmap='RdBu_r', vmin=-60, vmax=60)\n",
    "axes[0].set_title('ERA-5 reanalysis', fontsize=16, pad=10);\n",
    "axes[1].set_title('Observations', fontsize=16, pad=10);\n",
    "axes[2].set_title('Bias: ERA-5 - Observations', fontsize=16, pad=10)\n",
    "\n",
    "# adjust axes\n",
    "for ax in axes.flat:\n",
    "    ax.axes.get_xaxis().set_ticklabels([])\n",
    "    ax.axes.get_xaxis().set_ticks([])\n",
    "    ax.axes.get_yaxis().set_ticklabels([])\n",
    "    ax.axes.get_yaxis().set_ticks([])\n",
    "    ax.axes.axis('tight')\n",
    "    ax.set_xlabel('')\n",
    "    ax.set_ylabel('')\n",
    "    ax.set_axis_off()\n",
    "\n",
    "# adjust figure\n",
    "fig.suptitle('Average monthly precipitation (mm): 1980 - 2010', fontsize=20);\n",
    "fig.subplots_adjust(hspace=0, wspace=0, top=0.85)\n",
    "\n",
    "# add colorbar for bias\n",
    "axes = axes.flatten()\n",
    "cbar_ax_bias = fig.add_axes([axes[-1].get_position().x1 + 0.01, axes[-1].get_position().y0,\n",
    "                             0.01, axes[-1].get_position().y1 - axes[-1].get_position().y0])\n",
    "cbar_bias = fig.colorbar(im3, cax=cbar_ax_bias)\n",
    "cbar_bias.set_label(label='Relative bias / (%)', fontsize=16)\n",
    "cbar_bias.ax.tick_params(labelsize=14)\n",
    "\n",
    "# add colorbar for predictand\n",
    "cbar_ax_predictand = fig.add_axes([axes[0].get_position().x0, axes[0].get_position().y0 - 0.1,\n",
    "                                   axes[-1].get_position().x0 - axes[0].get_position().x0,\n",
    "                                   0.05])\n",
    "cbar_predictand = fig.colorbar(im1, cax=cbar_ax_predictand, orientation='horizontal')\n",
    "cbar_predictand.set_label(label='Precipitation / (mm month$^{-1}$)', fontsize=16)\n",
    "cbar_predictand.ax.tick_params(labelsize=14)\n",
    "\n",
    "# save figure\n",
    "fig.savefig('../Notebooks/Figures/capstone_pr.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc1a78dd-4912-4fc2-ab76-1a878f5e1b4c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}