Random_forest_check_depths.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d545b75d",
   "metadata": {},
   "source": [
    "### Comparing Random Forest performance for different depths"
   ]
  },
  {
   "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.ensemble import RandomForestClassifier, RandomForestRegressor\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",
    "from pathlib import Path\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8daf9b9-8e6b-44e3-9763-11942a521ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom utils\n",
    "from utils.utils_data import *\n",
    "from utils.utils_plot import *\n",
    "from utils.utils_RF import *\n",
    "#from utils.utils_ml 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",
    "PRECIP_DATA = 'ERA5-low'\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_XTRM = 0.95 # Percentile (threshold) for the extremes\n",
    "RF_MAX_DEPTH = [3,4,6,8,10] # testing different depths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98c274e-b72a-43c9-9ca0-be7cc16f1ca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Subsample domain for dev\n",
    "#LONS_INPUT = [10, 15]\n",
    "#LATS_INPUT = [40, 45]\n",
    "#LONS_PREC = [10, 15]\n",
    "#LATS_PREC = [40, 45]"
   ]
  },
  {
   "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",
    "pr = get_nc_data(PATH_ERA5 + '/precipitation/day_grid1/*nc', DATE_START, DATE_END, LONS_PREC, LATS_PREC)\n",
    "pr = pr.tp\n",
    "pr['time'] = pd.DatetimeIndex(pr.time.dt.date)"
   ]
  },
  {
   "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": "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c847649a",
   "metadata": {},
   "source": [
    "### Input data: meteorological fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce5176d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with dask.config.set(**{'array.slicing.split_large_chunks': False}):\n",
    "    # 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",
    "    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",
    "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": "3b82f636-3da7-4414-9619-0a57058b7c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge arrays\n",
    "with dask.config.set(**{'array.slicing.split_large_chunks': False}):\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",
    "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",
    "# we might not need to split into valid\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": "68f74e14-42d8-4e85-982a-21ca6f6f32be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_predictions(y_pred, dg, qq, pred_xtrm=False, show_plots=True, plot_most_extreme=True, plot_worst_best=False, plot_scores=True, plot_confusion_components=True):\n",
    "    if pred_xtrm:\n",
    "        print(\"exreme\")\n",
    "       # y_pred_bool = y_pred >= 0.5\n",
    "        y_pred_bool = y_pred  #no need\n",
    "    else:\n",
    "        print(\"non extremes\")\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",
    "    # Extract true values\n",
    "    \n",
    "    y_xtrm = dg.y_xtrm.to_numpy().squeeze()\n",
    "    y_prec = dg.y.to_numpy().squeeze()\n",
    "    \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",
    "    if show_plots and plot_most_extreme:\n",
    "        \n",
    "        if pred_xtrm:\n",
    "            print(\"plot\")\n",
    "            fig, axes = plt.subplots(figsize=(12, 3.5), ncols=2, nrows=1)\n",
    "            plot_map(axes[0], lons_y, lats_y, y_xtrm[i_max_obs], title=\"Day with max # extremes - truth (xtrm)\", vmin=0, vmax=1)\n",
    "            plot_map(axes[1], lons_y, lats_y, y_pred[i_max_obs], title=\"Day with max # extremes - prediction (prob xtrm)\", vmin=0, vmax=1)\n",
    "        else:\n",
    "            fig, axes = plt.subplots(figsize=(24, 3.5), ncols=4, nrows=1)\n",
    "            vmax = max(np.max(y_prec[i_max_obs]), np.max(y_pred[i_max_obs]))\n",
    "            plot_map(axes[0], lons_y, lats_y, y_prec[i_max_obs], title=\"Day with max # extremes - truth (val)\", vmin=0, vmax=vmax)\n",
    "            plot_map(axes[1], lons_y, lats_y, y_pred[i_max_obs], title=\"Day with max # extremes - prediction (val)\", vmin=0, vmax=vmax)\n",
    "            plot_map(axes[2], lons_y, lats_y, y_xtrm[i_max_obs], title=\"Day with max # extremes - truth (xtrm)\", vmin=0, vmax=1)\n",
    "            plot_map(axes[3], lons_y, lats_y, y_pred_bool[i_max_obs], title=\"Day with max # extremes - prediction (xtrm)\", vmin=0, vmax=1)\n",
    "        \n",
    "    # Get the index of the max/min difference between prediction and obs\n",
    "    y_diffs_series = np.sum(np.absolute(y_xtrm - y_pred_bool), axis=(1,2))\n",
    "   # i_worst_pred = np.argmax(y_diffs_series)\n",
    "    i_worst_pred = np.argmax(np.array(y_diffs_series)) # convert to array\n",
    "    i_best_pred = np.argmin(np.array(y_diffs_series))\n",
    "    \n",
    "    if show_plots and plot_worst_best:\n",
    "        \n",
    "        if pred_xtrm:\n",
    "            fig, axes = plt.subplots(figsize=(24, 3.5), ncols=4, nrows=1)\n",
    "            plot_map(axes[0], lons_y, lats_y, y_xtrm[i_worst_pred], title=\"Worst prediction - truth\", vmin=0, vmax=1)\n",
    "            plot_map(axes[1], lons_y, lats_y, y_pred[i_worst_pred], title=\"Worst prediction - prediction\", vmin=0, vmax=1)\n",
    "            plot_map(axes[2], lons_y, lats_y, y_xtrm[i_best_pred], title=\"Best prediction - truth\", vmin=0, vmax=1)\n",
    "            plot_map(axes[3], lons_y, lats_y, y_pred[i_best_pred], title=\"Best prediction - prediction\", vmin=0, vmax=1)\n",
    "        else:\n",
    "            fig, axes = plt.subplots(figsize=(24, 3.5), ncols=4, nrows=1)\n",
    "            plot_map(axes[0], lons_y, lats_y, y_xtrm[i_worst_pred], title=\"Worst prediction - truth\", vmin=0, vmax=1)\n",
    "            plot_map(axes[1], lons_y, lats_y, y_pred_bool[i_worst_pred], title=\"Worst prediction - prediction\", vmin=0, vmax=1)\n",
    "            plot_map(axes[2], lons_y, lats_y, y_xtrm[i_best_pred], title=\"Best prediction - truth\", vmin=0, vmax=1)\n",
    "            plot_map(axes[3], lons_y, lats_y, y_pred_bool[i_best_pred], title=\"Best prediction - prediction\", vmin=0, vmax=1)\n",
    "    \n",
    "    # Compute scores\n",
    "    \n",
    "    precision, recall = eval_confusion_matrix_scores_on_map(y_xtrm, y_pred_bool)\n",
    "    \n",
    "    if pred_xtrm:\n",
    "        print(\"get scores\")\n",
    "        roc_auc = eval_roc_auc_score_on_map(y_xtrm, y_pred)\n",
    "    if not pred_xtrm:\n",
    "        MSE = np.square(np.subtract(y_prec,y_pred)).mean() \n",
    "        rmse = sqrt(MSE)\n",
    "        \n",
    "    if show_plots and plot_scores:\n",
    "        if pred_xtrm:\n",
    "            fig, axes = plt.subplots(figsize=(18, 4), ncols=3, nrows=1)\n",
    "        else:\n",
    "            fig, axes = plt.subplots(figsize=(12, 4), ncols=2, nrows=1)\n",
    "\n",
    "        plot_map(axes[0], lons_y, lats_y, precision, title=\"Precision\", vmin=0, vmax=1)\n",
    "        plot_map(axes[1], lons_y, lats_y, recall, title=\"Recall\", vmin=0, vmax=1)\n",
    "        if pred_xtrm:\n",
    "            plot_map(axes[2], lons_y, lats_y, roc_auc, title=\"ROC AUC\", vmin=0.5, vmax=1)\n",
    "\n",
    "    if show_plots and plot_confusion_components:\n",
    "        tn, fp, fn, tp = eval_confusion_matrix_on_map(y_xtrm, y_pred_bool)\n",
    "        fig, axes = plt.subplots(figsize=(24, 4), ncols=4, nrows=1)\n",
    "        plot_map(axes[0], lons_y, lats_y, tn, title=\"True negative\")\n",
    "        plot_map(axes[1], lons_y, lats_y, fp, title=\"False positive\")\n",
    "        plot_map(axes[2], lons_y, lats_y, fn, title=\"False negative\")\n",
    "        plot_map(axes[3], lons_y, lats_y, tp, title=\"True positive\")\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "    if pred_xtrm:\n",
    "        print(\"return scores\")\n",
    "        return np.mean(precision), np.mean(recall), np.mean(roc_auc)\n",
    "    \n",
    "    return np.mean(precision), np.mean(recall), np.mean(rmse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "429afe9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "i_shape = dg_train.X.shape\n",
    "o_shape = dg_train.y.shape\n",
    "\n",
    "print(f'X shape: {i_shape}')\n",
    "print(f'y shape: {o_shape}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1538787e-2f13-4e04-9581-2ce873c62fc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk = {'lon':5, 'lat':5}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ba7bd3",
   "metadata": {},
   "source": [
    "# Train a random forest classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "039880c7-434a-4785-a0cc-529a42baadea",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_dask = dg_train.X.chunk(chunk)\n",
    "y_train_dask = dg_train.y_xtrm.chunk(chunk)\n",
    "\n",
    "# target for the regression\n",
    "yreg_train_dask = dg_train.y.chunk(chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6af28df-3078-4b48-957f-ed1f569089e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict the amount of precipitation\n",
    "df_prec = pd.DataFrame(columns = ['id', 'name',\n",
    "                                  'train_pr_rmse', 'test_pr_rmse', \n",
    "                                  'train_xtrm_precision', 'test_xtrm_precision', \n",
    "                                  'train_xtrm_recall', 'test_xtrm_recall'])\n",
    "\n",
    "df_xtrm = pd.DataFrame(columns = ['id', 'name', \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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c954b57-4844-402d-aa04-ab37e458742c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in RF_MAX_DEPTH:\n",
    "    \n",
    "    print('fit RF for depth',i)\n",
    "    # models file\n",
    "    clf_file = f'tmp/RF/trained_classifiers_RF_{PRECIP_DATA}_{PRECIP_XTRM}_{i}.pkl'\n",
    "    \n",
    "    if os.path.isfile(clf_file):\n",
    "        print('open RF for classification')\n",
    "        #open models\n",
    "        with open(clf_file, 'rb') as f:\n",
    "            crfs = pickle.load(f)\n",
    "    else:\n",
    "\n",
    "        crfs = xr.apply_ufunc(\n",
    "            train_rf_classifier_model,\n",
    "            X_train_dask, y_train_dask, i,\n",
    "            vectorize=True,\n",
    "            dask = 'parallelized',\n",
    "            input_core_dims=[['time', 'level'], ['time'], []],  # reduce along these dimensions\n",
    "            output_dtypes=[object]\n",
    "        ).compute()\n",
    "        # save the models\n",
    "        with open(clf_file, 'wb') as output:\n",
    "                pickle.dump(crfs, output, pickle.HIGHEST_PROTOCOL)\n",
    "                \n",
    "    # calculate scores\n",
    "    print('predicting test')\n",
    "    y_test_preds = xr.apply_ufunc(\n",
    "                    apply_rf_classifier_model, \n",
    "                    crfs, dg_test.X,\n",
    "                    vectorize=True,\n",
    "                    input_core_dims=[[],['time', 'level']],\n",
    "                    output_dtypes=['object'] # change the output type\n",
    "                    ).compute()\n",
    "    \n",
    "    y_test_preds_proba = xr.apply_ufunc(\n",
    "                        apply_rf_classifier_model_proba, \n",
    "                        crfs, dg_test.X,\n",
    "                        vectorize=True,\n",
    "                        input_core_dims=[[],['time', 'level']],\n",
    "                        output_dtypes=['object'] # change the output type\n",
    "                    ).compute()\n",
    "    \n",
    "    print('predicting train')\n",
    "    y_train_preds = xr.apply_ufunc(\n",
    "                    apply_rf_classifier_model, \n",
    "                    crfs, dg_train.X,\n",
    "                    vectorize=True,\n",
    "                    input_core_dims=[[],['time', 'level']],\n",
    "                    output_dtypes=['object'] # change the output type\n",
    "                ).compute()\n",
    "    \n",
    "    y_pred_xtrm_test = create_xarray_frompred(y_test_preds, dg_test.y.shape[0], lats_y, lons_x)\n",
    "    yprob_pred_xtrm_test = create_xarray_frompred(y_test_preds_proba,  dg_test.y.shape[0], lats_y, lons_x)\n",
    "    y_pred_xtrm_train = create_xarray_frompred(y_train_preds,  dg_train.y.shape[0], lats_y, lons_x)\n",
    "    \n",
    "    print('Train scores')\n",
    "    precision_train, recall_train, roc_auc_train = analyze_predictions(y_pred_xtrm_train, dg_train, qq, pred_xtrm=True, show_plots=False)\n",
    "    \n",
    "    \n",
    "    df_xtrm = df_xtrm.append({'id': i,  'name': 'Classifier'}, ignore_index=True)\n",
    "    \n",
    "    df_xtrm.at[df_xtrm.index[-1], 'name'] = 'Classifier'\n",
    "    df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_precision'] = precision_train\n",
    "    df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_recall'] = recall_train\n",
    "    df_xtrm.at[df_xtrm.index[-1], 'train_xtrm_roc_auc'] = roc_auc_train\n",
    "    print('Test scores')\n",
    "    precision, recall, roc_auc_test = analyze_predictions(y_pred_xtrm_test, dg_test, qq, pred_xtrm=False, 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_test\n",
    "            \n",
    "    \n",
    "    # Regressor part\n",
    "    # models file\n",
    "    rfs_file = f'tmp/RF/trained_regressors_RF_{PRECIP_DATA}_{PRECIP_XTRM}_{i}.pkl'\n",
    "    \n",
    "\n",
    "    if os.path.isfile(rfs_file):\n",
    "        #print('open models')\n",
    "        #open models\n",
    "        with open(rfs_file, 'rb') as f:\n",
    "            rfs = pickle.load(f)\n",
    "    else:\n",
    "\n",
    "        rfs = xr.apply_ufunc(\n",
    "            train_rf_regress_model,\n",
    "            X_train_dask, yreg_train_dask, i,\n",
    "            vectorize=True,\n",
    "            dask = 'parallelized',\n",
    "            input_core_dims=[['time', 'level'], ['time'], []],  # reduce along these dimensions\n",
    "            output_dtypes=[object]\n",
    "        ).compute()\n",
    "        # save the models\n",
    "        with open(rfs_file, 'wb') as output:\n",
    "                pickle.dump(rfs, output, pickle.HIGHEST_PROTOCOL)\n",
    "                \n",
    "                \n",
    "                \n",
    "    # calculate predictions\n",
    "    y_test_preds_reg = xr.apply_ufunc(\n",
    "        apply_rf_regress_model, \n",
    "        rfs, dg_test.X,\n",
    "        vectorize=True,\n",
    "        input_core_dims=[[],['time', 'level']],\n",
    "        output_core_dims=[[]],\n",
    "        output_dtypes=[object] # change the output type\n",
    "    ).compute()\n",
    "        \n",
    "        # calculate predictions\n",
    "    y_train_preds_reg = xr.apply_ufunc(\n",
    "        apply_rf_regress_model, \n",
    "        rfs, dg_train.X,\n",
    "        vectorize=True,\n",
    "        input_core_dims=[[],['time', 'level']],\n",
    "        output_core_dims=[[]],\n",
    "        output_dtypes=[object] # change the output type\n",
    "    ).compute()\n",
    "        \n",
    "        \n",
    "    # we need to create matrix \n",
    "    preds_test_reg_matrix = create_xarray_frompred(y_test_preds_reg, dg_test.y.shape[0], lats_y, lons_x)\n",
    "    preds_train_reg_matrix = create_xarray_frompred(y_train_preds_reg, dg_train.y.shape[0], lats_y, lons_x)\n",
    "\n",
    "    print('Train results regressor')\n",
    "    \n",
    "    precision_train, recall_train, rmse_train = analyze_predictions(preds_train_reg_matrix, dg_train, qq, pred_xtrm=False, show_plots=False)\n",
    "       \n",
    "    df_prec = df_prec.append({'id': i,  'name': 'Regressor'}, ignore_index=True)\n",
    "        \n",
    "    df_prec.at[df_prec.index[-1], 'name'] = 'Regressor'\n",
    "    df_prec.at[df_prec.index[-1], 'train_xtrm_precision'] = precision_train\n",
    "    df_prec.at[df_prec.index[-1], 'train_xtrm_recall'] = recall_train\n",
    "    df_prec.at[df_prec.index[-1], 'train_pr_rmse'] = rmse_train\n",
    "    print('Test results regressor')\n",
    "    precision, recall, rmse = analyze_predictions(preds_test_reg_matrix, dg_test, qq, pred_xtrm=False, show_plots=False)\n",
    "        \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",
    "    df_prec.at[df_prec.index[-1], 'test_pr_rmse'] = rmse\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a791b07-567c-4e86-b47c-4263a2d91856",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_prec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07607203-fcf5-46b9-a9ef-7212f6246c79",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_xtrm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1000a3e8-0786-428a-af33-328fab693fd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_prec.to_csv('tmp/df_prec_RF_depth.csv')\n",
    "df_xtrm.to_csv('tmp/df_xtrm_RF_depth.csv')"
   ]
  }
 ],
 "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
}