2_Evaluate_fit_LGBM_model_TSCV.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate and fit a ML model on the EC flux tower data\n",
    "\n",
    "This notebook is largely for testing/experimenting with different model configurations. But it also allows for exploring feature importance and dependencies between variables (which were used the manuscript). The model output here is not used in the final predictions, the models output in the `3_Generate_ensembles.ipynb` notebook are used instead.\n",
    "\n",
    "Using a custom, nested **time-series cross validation** approach for model evaulation.\n",
    "\n",
    "Notes on experimental model iterations:\n",
    "\n",
    "Testing features: beginning with these:\n",
    "\n",
    "     'kNDVI_anom_RS',\n",
    "     'FPAR_RS',\n",
    "     'LST_RS',\n",
    "     'tree_cover_RS',\n",
    "     'nontree_cover_RS',\n",
    "     'nonveg_cover_RS',\n",
    "     'LST-Tair_RS',\n",
    "     'TWI_RS',\n",
    "     'NDWI_RS',\n",
    "     'rain_anom_RS',\n",
    "     'rain_cml3_anom_RS',\n",
    "     'rain_cml6_anom_RS',\n",
    "     'rain_cml12_anom_RS',\n",
    "     'srad_anom_RS',\n",
    "     'vpd_RS',\n",
    "     'tavg_anom_RS',\n",
    "     'SOC_RS',\n",
    "     'C4_percent_RS'\n",
    "\n",
    "PLUS:\n",
    "\n",
    "- `20221122` inlcuded MI, Elevation, and VegH (where MI using pan evaporation)\n",
    "- `20221123` included MI, Elevation (where MI using pan evaporation)\n",
    "- `20221128` included MI (where MI using pan evaporation)\n",
    "- `20221129` All of the above features removed.\n",
    "- `20221130` Including MI but where MI uses AWRA Penman–Monteith PET instead of pan evap.\n",
    "- `20221206` Including Penman–Monteith MI, removed Daly Uncleared, added kNDVI, update EC data with 2022 data\n",
    "- `20221208` As per 20221206 but using fPAR-NDVI (Donohue) instead of MODIS fPAR\n",
    "- `20221212` Remove fPAR. Update climate vars to use actual values rather than anomalies (srad, tavg, rain) (keep rainfall anomalies)\n",
    "- `20230109` Remove monthly rain_anom (already contained in 3-month cml_anom), add back srad_anom and tavg_anom. Remove MI (can't run at 1km). TiTreeEast site added (coords fixed)\n",
    "- ~~`20230112` Remove cumulative rainfall anomalies and replace with just cumulative rainfall.~~\n",
    "- `20230306` Same as 20230109 but remove SOC and add rain_anom\n",
    "- `20230316` Replacing MODIS VCF with Randall's Trees/Grass/Bare + C4_grass and reintroducing Veg Height (VegH)\n",
    "- `20230320` Same as 20230316 but removing TWI as it creates artefacts when its feature importance goes up.\n",
    "\n",
    "Models without satellite data:\n",
    "\n",
    "Testing features: \n",
    "\n",
    "     'tree_cover_RS',\n",
    "     'nontree_cover_RS',\n",
    "     'nonveg_cover_RS',\n",
    "     'TWI_RS',\n",
    "     'rain_RS',\n",
    "     'rain_cml3_anom_RS',\n",
    "     'rain_cml6_anom_RS',\n",
    "     'rain_cml12_anom_RS',\n",
    "     'srad_RS',\n",
    "     'srad_anom_RS',\n",
    "     'vpd_RS',\n",
    "     'tavg_RS',\n",
    "     'tavg_anom_RS',\n",
    "     'SOC_RS',\n",
    "     'C4_percent_RS',\n",
    "     'month_RS',\n",
    "\n",
    "- `CLIM_20221223`\n",
    "- `CLIM_20230112` #same as above but including update to TD\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sb\n",
    "from scipy import stats\n",
    "from joblib import dump\n",
    "from pprint import pprint\n",
    "from scipy.stats import gaussian_kde\n",
    "import matplotlib.pyplot as plt\n",
    "import lightgbm as lgbm\n",
    "from lightgbm import LGBMRegressor\n",
    "\n",
    "import shap\n",
    "from sklearn.model_selection import TimeSeriesSplit, RandomizedSearchCV, KFold\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'AUS'\n",
    "model_var = 'GPP'\n",
    "suffix = '20230320'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base = '/g/data/os22/chad_tmp/NEE_modelling/results/training_data/'\n",
    "sites = os.listdir('/g/data/os22/chad_tmp/NEE_modelling/results/training_data/')\n",
    "\n",
    "td = []\n",
    "for site in sites:\n",
    "    if '.csv' in site:\n",
    "        xx = pd.read_csv(base+site, index_col='time', parse_dates=True)\n",
    "        xx['site'] = site[0:5]\n",
    "        td.append(xx)\n",
    "\n",
    "ts = pd.concat(td).dropna() #we'll use this later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variables = [\n",
    "             'kNDVI_RS',\n",
    "             'kNDVI_anom_RS',\n",
    "             'LST_RS',\n",
    "             'trees_RS',\n",
    "             'grass_RS',\n",
    "             'bare_RS',\n",
    "             'C4_grass_RS',\n",
    "             'LST-Tair_RS',\n",
    "             'NDWI_RS',\n",
    "             'rain_RS',\n",
    "             'rain_anom_RS',\n",
    "             'rain_cml3_anom_RS',\n",
    "             'rain_cml6_anom_RS',\n",
    "             'rain_cml12_anom_RS',\n",
    "             'srad_RS',\n",
    "             'srad_anom_RS',\n",
    "             'vpd_RS',\n",
    "             'tavg_RS',\n",
    "             'tavg_anom_RS',\n",
    "             'VegH_RS',\n",
    "             'site'\n",
    "            ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xx = []\n",
    "yy = []\n",
    "\n",
    "for t in td:\n",
    "    # t = t.drop(['Fluxcom_RS-Meteo_NEE', 'Fluxcom_RS_NEE', 'ThisStudy_NEE', 'Cable_NEE',\n",
    "    #    'Fluxcom_RS_GPP', 'Fluxcom_RS-meteo_GPP', 'ThisStudy_GPP', 'Cable_GPP',\n",
    "    #    'MODIS_GPP', 'GOSIF_GPP'], axis=1)  \n",
    "    \n",
    "    t = t.dropna()  # remove NaNS\n",
    "    df = t.drop(['NEE_SOLO_EC','GPP_SOLO_EC','ER_SOLO_EC'], axis=1) # seperate carbon fluxes\n",
    "    \n",
    "    #df = df.filter(regex='RS') # only use remote sensing variables   \n",
    "    df = df[variables]\n",
    "    \n",
    "    if model_var == 'ET':\n",
    "        df_var=t[model_var+'_EC']\n",
    "    else:\n",
    "        df_var=t[model_var+'_SOLO_EC'] # seperate out the variable we're modelling\n",
    "    \n",
    "    x = df.reset_index(drop=True)#.to_numpy()\n",
    "    y = df_var.reset_index(drop=True)#.to_numpy()\n",
    "    xx.append(x)\n",
    "    yy.append(y)\n",
    "\n",
    "x = pd.concat(xx)\n",
    "y = pd.concat(yy)\n",
    "\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test model robustness with time-series K-fold cross validation\n",
    "\n",
    "* If you set boosting as RF then the lightgbm algorithm behaves as random forest. According to the documentation, to use RF you must use bagging_fraction and feature_fraction smaller than 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate five sets of train-test indices\n",
    "\n",
    "For each site, grab a sequential set of test samples (time-series-split methods), the remaining points (either side of test samples) go into training.  A single K-fold contains test and training samples from every site."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sites = x['site'].unique()\n",
    "x['original_index'] = [i for i in range(0,len(x))]\n",
    "\n",
    "train_1=[]\n",
    "train_2=[]\n",
    "train_3=[]\n",
    "train_4=[]\n",
    "train_5=[]\n",
    "\n",
    "test_1=[]\n",
    "test_2=[]\n",
    "test_3=[]\n",
    "test_4=[]\n",
    "test_5=[]\n",
    "\n",
    "for site in sites:\n",
    "    df = x.loc[x['site'] == site]\n",
    "    tscv = TimeSeriesSplit(n_splits=5)\n",
    "    i=1\n",
    "    for train, test in tscv.split(df):\n",
    "        all_indices=np.concatenate([train,test])\n",
    "        left_over = df.loc[~df.index.isin(all_indices)].index.values\n",
    "        train = np.concatenate([train, left_over])\n",
    "        if i==1:\n",
    "            train_1.append(df.iloc[train]['original_index'].values)\n",
    "            test_1.append(df.iloc[test]['original_index'].values)\n",
    "        if i==2:\n",
    "            train_2.append(df.iloc[train]['original_index'].values)\n",
    "            test_2.append(df.iloc[test]['original_index'].values)\n",
    "        if i==3:\n",
    "            train_3.append(df.iloc[train]['original_index'].values)\n",
    "            test_3.append(df.iloc[test]['original_index'].values)\n",
    "        if i==4:\n",
    "            train_4.append(df.iloc[train]['original_index'].values)\n",
    "            test_4.append(df.iloc[test]['original_index'].values)\n",
    "        if i==4:\n",
    "            train_5.append(df.iloc[train]['original_index'].values)\n",
    "            test_5.append(df.iloc[test]['original_index'].values)\n",
    "        i+=1\n",
    "\n",
    "train_1 = np.concatenate(train_1)\n",
    "train_2 = np.concatenate(train_2)\n",
    "train_3 = np.concatenate(train_3)\n",
    "train_4 = np.concatenate(train_4)\n",
    "train_5 = np.concatenate(train_5)\n",
    "\n",
    "test_1 = np.concatenate(test_1)\n",
    "test_2 = np.concatenate(test_2)\n",
    "test_3 = np.concatenate(test_3)\n",
    "test_4 = np.concatenate(test_4)\n",
    "test_5 = np.concatenate(test_5)\n",
    "\n",
    "train = [train_1, train_2, train_3, train_4, train_5]\n",
    "test = [test_1, test_2, test_3, test_4, test_5]\n",
    "\n",
    "#check there are no train indices in the test indices\n",
    "for i,j in zip(train, test):\n",
    "    assert (np.sum(np.isin(i,j)) == 0)\n",
    "\n",
    "#remove the columns we no longer need\n",
    "x = x.drop(['site', 'original_index'], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Consider using SHAP-based RFA feature selection\n",
    "\n",
    "Ulimately, feature selection was not used as each method gives confliting results, and we don't have too many features anyway.\n",
    "\n",
    "<!-- python library: https://github.com/cerlymarco/shap-hypetune\n",
    "\n",
    "Recursive feature addition: https://towardsdatascience.com/recursive-feature-selection-addition-or-elimination-755e5d86a791 -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the parameter grid using distributions\n",
    "param_grid = {\n",
    "    'num_leaves': stats.randint(5,50),\n",
    "    'min_child_samples':stats.randint(10,30),\n",
    "    'boosting_type': ['gbdt', 'dart'],\n",
    "    'max_depth': stats.randint(5,25),\n",
    "    'n_estimators': [200, 300, 400, 500],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# clf = BoostRFA(  \n",
    "#     LGBMRegressor(random_state=0, n_jobs=-1),                        \n",
    "#     # min_features_to_select=10,              # the minimum number of features to be selected  \n",
    "#     step=1,                                 # number of features to remove at each iteration  \n",
    "#     param_grid=param_grid,                  # parameters to be optimized  \n",
    "#     greater_is_better=False,                # minimize or maximize the monitored score  \n",
    "#     importance_type='shap_importances',     # which importance measure to use: features or shap  \n",
    "#     train_importance=False,                 # False=use eval_set  \n",
    "#     n_iter=10,                              # number of sampled parameter configurations   \n",
    "#     sampling_seed=1,\n",
    "#     verbose=0,                                \n",
    "#     n_jobs=-1                              \n",
    "# )  \n",
    "\n",
    "# #pick one fold to use as the validation/training split\n",
    "# X_tr, X_tt = x.iloc[train[0], :], x.iloc[test[0], :]\n",
    "# y_tr, y_tt = y.iloc[train[0]], y.iloc[test[0]]\n",
    "\n",
    "# #test models\n",
    "# clf.fit(X_tr, y_tr,\n",
    "#         eval_set=[(X_tt, y_tt)],\n",
    "#         # early_stopping_rounds=5,\n",
    "#         verbose=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print('Num of features selected: ', clf.n_features_)\n",
    "# print('Num of features selected: ', clf.best_params_)\n",
    "# print('Features selected:', X_tr[X_tr.columns[clf.support_]].columns.values)\n",
    "\n",
    "#write out selected features for each fold\n",
    "# textfile = open(\"/g/data/os22/chad_tmp/NEE_modelling/results/features_select_\"+var+\".txt\", \"w\")\n",
    "# for element in x[X_tr.columns[clf.support_]].columns.values:\n",
    "#     textfile.write(element + \",\")\n",
    "# textfile.close()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write out text file of feature layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "textfile = open(\"/g/data/os22/chad_tmp/NEE_modelling/results/variables_\"+suffix+\".txt\", \"w\")\n",
    "for element in x.columns:\n",
    "    textfile.write(element + \",\")\n",
    "textfile.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remove features not selected (if using)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x = x[x.columns[clf.support_]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conduct nested cross-validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lists to store results of CV testing\n",
    "acc = []\n",
    "rmse=[]\n",
    "r2=[]\n",
    "\n",
    "i = 1\n",
    "for train_index, test_index in zip(train, test):\n",
    "    print(f\"Working on {i}/{len(train)} outer cv split\", end='\\r')\n",
    "    model = LGBMRegressor(random_state=1,\n",
    "                          verbose=-1,\n",
    "                          # objective='mean_absolute_percentage_error'\n",
    "                          # monotone_constraints=m_con,\n",
    "                          # monotone_constraints_method='intermediate'\n",
    "                          )\n",
    "\n",
    "    # index training, testing\n",
    "    X_tr, X_tt = x.iloc[train_index, :], x.iloc[test_index, :]\n",
    "    y_tr, y_tt = y.iloc[train_index], y.iloc[test_index]\n",
    "    \n",
    "    #simple random split on inner fold\n",
    "    inner_cv = KFold(n_splits=5,\n",
    "                     shuffle=True,\n",
    "                     random_state=0)\n",
    "    \n",
    "    clf = RandomizedSearchCV(\n",
    "                   model,\n",
    "                   param_grid,\n",
    "                   verbose=0,\n",
    "                   n_iter=250,\n",
    "                   n_jobs=-1,\n",
    "                   cv=inner_cv.split(X_tr, y_tr)\n",
    "                  )\n",
    "    \n",
    "    #prevents extensive print statements\n",
    "    clf.fit(X_tr, y_tr, callbacks=None)\n",
    "    \n",
    "    # predict using the best model\n",
    "    best_model = clf.best_estimator_\n",
    "    pred = best_model.predict(X_tt)\n",
    "\n",
    "    # evaluate model w/ multiple metrics\n",
    "    # r2\n",
    "    r2_ = r2_score(y_tt, pred)\n",
    "    r2.append(r2_)\n",
    "    # MAE\n",
    "    ac = mean_absolute_error(y_tt, pred)\n",
    "    acc.append(ac)\n",
    "    # RMSE\n",
    "    rmse_ = np.sqrt(mean_squared_error(y_tt, pred))\n",
    "    rmse.append(rmse_)\n",
    "    \n",
    "    #1:1 plots for each fold (save to csv so we can make a plot later on)\n",
    "    df = pd.DataFrame({'Test':y_tt, 'Pred':pred}).reset_index(drop=True)\n",
    "\n",
    "    df.to_csv(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/\"+str(i)+\"_\"+model_var+\"_lgbm.csv\")\n",
    "    \n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Print accuracy results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Mean MAE accuracy: \"+ str(round(np.mean(acc), 2)))\n",
    "print(\"Std dev of MAE accuracy: \"+ str(round(np.std(acc), 2)))\n",
    "print('\\n')\n",
    "print(\"Mean RMSE: \"+ str(round(np.mean(rmse), 2)))\n",
    "print(\"Std dev RMSE: \"+ str(round(np.std(rmse), 2)))\n",
    "print('\\n')\n",
    "print(\"Mean r2: \"+ str(round(np.mean(r2), 2)))\n",
    "print(\"Std dev r2: \"+ str(round(np.std(r2), 2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a single 1:1 plot out of the folds \n",
    "\n",
    "None of the test samples overlap between folds, and every sample has been tested\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dffs=[]\n",
    "for i in range(1,len(train)+1):\n",
    "    df = pd.read_csv(f'/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/{model_var}_ensemble/CV{i}_{model_var}_lgbm.csv', usecols=['Test', 'Pred'])\n",
    "    dffs.append(df)\n",
    "\n",
    "cross_df = pd.concat(dffs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,1, figsize=(7,6))\n",
    "\n",
    "xy = np.vstack([cross_df['Test'],cross_df['Pred']])\n",
    "z = gaussian_kde(xy)(xy)\n",
    "\n",
    "sb.scatterplot(data=cross_df, x='Test',y='Pred',c=z, s=50, lw=1, alpha=0.5, ax=ax)\n",
    "sb.regplot(data=cross_df, x='Test',y='Pred', scatter=False, color='darkblue', ax=ax)\n",
    "sb.regplot(data=cross_df, x='Test',y='Test', color='black', scatter=False, line_kws={'linestyle':'dashed'}, ax=ax);\n",
    "\n",
    "plt.xlabel('Observation '+ model_var + ' gC/m\\N{SUPERSCRIPT TWO}/month', fontsize=20)\n",
    "plt.ylabel('Prediction ' + model_var+ ' gC/m\\N{SUPERSCRIPT TWO}/month', fontsize=20)\n",
    "ax.text(.05, .95, 'r\\N{SUPERSCRIPT TWO}={:.2f}'.format(np.mean(r2)),\n",
    "            transform=ax.transAxes, fontsize=20)\n",
    "ax.text(.05, .9, 'MAE={:.3g}'.format(np.mean(acc)),\n",
    "            transform=ax.transAxes, fontsize=20)\n",
    "ax.tick_params(axis='x', labelsize=20)\n",
    "ax.tick_params(axis='y', labelsize=20)\n",
    "\n",
    "plt.tight_layout()\n",
    "fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/cross_val_\"+model_var+\"_lgbm_\"+suffix+\".png\",\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimize model using all training data\n",
    "\n",
    "Using a randomized strategy so we can search through more variables, with 500 iterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_grid = {\n",
    "    'num_leaves': stats.randint(5,50),\n",
    "    'min_child_samples':stats.randint(10,30),\n",
    "    'boosting_type': ['gbdt', 'dart'],\n",
    "    'max_depth': stats.randint(5,25),\n",
    "    'n_estimators': [300, 400, 500],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#generate n_splits of train-test_split\n",
    "#rs = ShuffleSplit(n_splits=outer_cv_splits, test_size=test_size, random_state=1)\n",
    "\n",
    "# instatiate a gridsearchCV\n",
    "# clf = GridSearchCV(LGBMRegressor(verbose=-1\n",
    "#                                  # boosting_type='rf',\n",
    "#                                  # bagging_freq=1, \n",
    "#                                  # bagging_fraction=0.8\n",
    "#                                 ),\n",
    "#                    param_grid,\n",
    "#                    scoring='r2',\n",
    "#                    verbose=1,\n",
    "#                    cv=zip(train, test), #using timeseries custom splits here\n",
    "#                   )\n",
    "\n",
    "clf = RandomizedSearchCV(LGBMRegressor(verbose=-1),\n",
    "                   param_grid,\n",
    "                   verbose=1,\n",
    "                   n_iter=500,\n",
    "                   n_jobs=-1,\n",
    "                   cv=zip(train, test), #using timeseries custom splits here\n",
    "                  )\n",
    "\n",
    "clf.fit(x, y, callbacks=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"The most accurate combination of tested parameters is: \")\n",
    "pprint(clf.best_params_)\n",
    "print('\\n')\n",
    "print(\"The best score using these parameters is: \")\n",
    "print(round(clf.best_score_, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit on all data using best params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LGBMRegressor(**clf.best_params_)\n",
    "\n",
    "model.fit(x,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump(model, '/g/data/os22/chad_tmp/NEE_modelling/results/models/'+model_name+'_'+model_var+'_LGBM_model_'+suffix+'.joblib')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine feature importance using SHAP\n",
    "\n",
    "SHAP (SHapley Additive exPlanations) is a game theoretic approach to explain the output of any machine learning model\n",
    "\n",
    "https://github.com/slundberg/shap\n",
    "\n",
    "If loading a previously saved model, uncomment the below cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import load\n",
    "model_var = 'NEE'\n",
    "suffix = '20230320'\n",
    "model_path = '/g/data/os22/chad_tmp/NEE_modelling/results/models/AUS_'+model_var+'_LGBM_model_'+suffix+'.joblib'\n",
    "model = load(model_path).set_params(n_jobs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# explain the model's predictions using SHAP\n",
    "explainer = shap.Explainer(model)\n",
    "shap_values = explainer(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals= np.abs(shap_values.values).mean(0)\n",
    "feature_importance = pd.DataFrame(list(zip(x.columns, vals)), columns=['col_name','feature_importance_vals'])\n",
    "feature_importance.sort_values(by=['feature_importance_vals'],ascending=False,inplace=True)\n",
    "feature_importance['col_name'] = feature_importance['col_name'].str.removesuffix(\"_RS\")\n",
    "# feature_importance.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(5,11))\n",
    "shap.summary_plot(shap_values, max_display=10, show=False, feature_names=feature_importance['col_name'])\n",
    "plt.gcf().axes[-1].set_aspect('auto')\n",
    "plt.gcf().axes[-1].set_box_aspect(15) \n",
    "ax.tick_params(axis='x', labelsize=20)\n",
    "ax.tick_params(axis='y', labelsize=20)\n",
    "ax.set_xlabel(model_var+' SHAP Value', fontsize=20)\n",
    "plt.tight_layout()\n",
    "# fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/feature_importance_\"+model_var+\"_lgbm_\"+suffix+\".png\", bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependency scatter plots (including main interaction effect)\n",
    "\n",
    "showing the effect a single features have on the predictions made by the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First plots show the dependence for the top 4 predictors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "font=28"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,4, figsize=(32,7))\n",
    "\n",
    "for ax, feature in zip(axs.ravel(), list(feature_importance['col_name'])[0:4]):\n",
    "    shap.plots.scatter(shap_values[:,feature+'_RS'], ax=ax, show=False, color=shap_values, cmap='viridis')\n",
    "    ax.set_ylabel('SHAP value '+ feature,  fontsize=font)\n",
    "    ax.set_xlabel(feature, fontsize=font)\n",
    "    ax.tick_params(axis='x', labelsize=font)\n",
    "    ax.tick_params(axis='y', labelsize=font)\n",
    "    plt.gcf().axes[-2].set_aspect('auto')\n",
    "    plt.gcf().axes[-2].set_box_aspect(20)\n",
    "    plt.gcf().axes[-2].yaxis.label.set_fontsize(font)\n",
    "    plt.gcf().axes[-2].tick_params(axis='y', labelsize=font)\n",
    "plt.tight_layout()\n",
    "fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/interaction_scatterplots_\"+model_var+\"_lgbm_\"+suffix+\".png\", bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second plots show the dependence for the four main climate predictors (with kNDVI interaction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,5, figsize=(40,7))\n",
    "\n",
    "for ax, feature, units in zip(axs.ravel(), ['kNDVI', 'srad', 'tavg', 'rain', 'vpd'], ['', 'MJ m⁻\\N{SUPERSCRIPT TWO} day⁻¹', '($^\\circ$C)', 'mm/month', 'HPa']): #list(feature_importance['col_name'][0:1])+\n",
    "    if feature =='kNDVI':\n",
    "        shap.plots.scatter(shap_values[:,feature+'_RS'], ax=ax, show=False, color=shap_values, cmap='viridis')\n",
    "    else:\n",
    "        shap.plots.scatter(shap_values[:,feature+'_RS'], ax=ax, show=False, color=shap_values[:,'kNDVI_RS'], cmap='viridis')\n",
    "        plt.gcf().axes[-2].set_ylabel('kNDVI')\n",
    "    ax.set_ylabel(feature+' SHAP values',  fontsize=font)\n",
    "    ax.set_xlabel(feature+', '+units+'', fontsize=font)\n",
    "    ax.tick_params(axis='x', labelsize=font)\n",
    "    ax.tick_params(axis='y', labelsize=font)\n",
    "    plt.gcf().axes[-2].set_aspect('auto')\n",
    "    plt.gcf().axes[-2].set_box_aspect(20)\n",
    "    plt.gcf().axes[-2].yaxis.label.set_fontsize(font)\n",
    "    plt.gcf().axes[-2].tick_params(axis='y', labelsize=font)\n",
    "plt.tight_layout()\n",
    "fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/interaction_scatterplots_climatevars_\"+model_var+\"_lgbm_\"+suffix+\".png\", bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}