3_Generate_ensemble_of_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate an ensemble of ML models\n",
    "\n",
    "We will attempt to model a portion of the empirical uncertainty that comes from the training data and optimization algorithm.  To do this, we will generate 30 models. For each iteration, two flux tower sites entire time-series will be removed from the training data and both a LGBM and RF will be fit on the remaining data.  This will result in 30 models that later we can use to make predictions 30 predictions. The IQR envelope of our predictions will inform our uncertainity\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 seaborn as sb\n",
    "import pandas as pd\n",
    "from joblib import dump\n",
    "from scipy import stats\n",
    "import matplotlib.pyplot as plt\n",
    "from pprint import pprint\n",
    "from scipy.stats import gaussian_kde\n",
    "from lightgbm import LGBMRegressor\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "import lightgbm as lgbm\n",
    "from sklearn.model_selection import TimeSeriesSplit, RandomizedSearchCV, KFold\n",
    "from sklearn.metrics import mean_absolute_error\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_var = 'NEE'\n",
    "suffix = '20230320'\n",
    "features_list = '/g/data/os22/chad_tmp/NEE_modelling/results/variables_'+suffix+'.txt'"
   ]
  },
  {
   "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": [
    "train_vars = list(pd.read_csv(features_list))[0:-1]\n",
    "train_vars=[i[:-3] for i in train_vars]\n",
    "train_vars = [i+'_RS' for i in train_vars]\n",
    "train_vars.append('site')\n",
    "train_vars"
   ]
  },
  {
   "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', 'ET_EC'], axis=1) # seperate fluxes\n",
    "    \n",
    "    #df = df.filter(regex='RS') # only use remote sensing variables   \n",
    "    df = df[train_vars]\n",
    "    \n",
    "    if model_var == 'ET':\n",
    "        df_var=t[[model_var+'_EC', 'site']]\n",
    "    else:\n",
    "        df_var=t[[model_var+'_SOLO_EC', 'site']] # 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)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate 30 iterations of models\n",
    "\n",
    "For two regression methods (RF and LigtGBM), we remove two randomly selected sites from the training data\n",
    "\n",
    "Then, we do the per site TSCV: 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.\n",
    "\n",
    "A model is built and saved that is trained of 15 iterations of site removal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "i=0\n",
    "for m in range(1,16): # 15 iterations for each model\n",
    "    print(\" {:03}/{:03}\\r\".format(m, len(range(1,16))))\n",
    "          \n",
    "    #randomly select two sites to remove from dataset\n",
    "    subset=np.random.choice(x['site'].unique(), size=2)\n",
    "    x_n = x[~x.site.isin(subset)]\n",
    "    y_n = y[~y.site.isin(subset)]\n",
    "\n",
    "    sites_n = x_n['site'].unique()\n",
    "    x_n['original_index'] = [i for i in range(0,len(x_n))]\n",
    "    \n",
    "    #build TSCV splits across all remaining sites\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_n in sites_n:\n",
    "        df = x_n.loc[x_n['site'] == site_n]\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_n = x_n.drop(['site', 'original_index'], axis=1)\n",
    "    y_n = y_n.drop('site', axis=1)\n",
    "\n",
    "    #loop through the two regression methods\n",
    "    for regressor in [LGBMRegressor,\n",
    "                      RandomForestRegressor]:\n",
    "\n",
    "        if isinstance(regressor(), lgbm.sklearn.LGBMRegressor):\n",
    "            m_name='_lgbm_'\n",
    "\n",
    "            param_grid_lgbm = {\n",
    "                'num_leaves': stats.randint(5,40),\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",
    "            }\n",
    "\n",
    "        else:\n",
    "            m_name='_rf_'\n",
    "\n",
    "            param_grid_rf = {\n",
    "                'max_depth': stats.randint(5,35),\n",
    "                'max_features': ['log2', None, \"sqrt\"],\n",
    "                'n_estimators': [200,300,400,500]}\n",
    "\n",
    "        print('  Model:', m_name)\n",
    "        \n",
    "        #-----Nested CV to test accuracy-----------------------------------------------\n",
    "        # results are saved as a .csv \n",
    "        j=1\n",
    "        for train_index, test_index in zip(train, test):\n",
    "            print(f\"    {j}/{len(train)} outer cv split\")\n",
    "            \n",
    "            if os.path.exists(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/\"+model_var+\"_ensemble/CV_\"+str(j)+\"_\"+model_var+m_name+str(m)+\".csv\"):\n",
    "                j+=1\n",
    "                continue\n",
    "\n",
    "            # index training, testing\n",
    "            X_tr, X_tt = x_n.iloc[train_index, :], x_n.iloc[test_index, :]\n",
    "            y_tr, y_tt = y_n.iloc[train_index], y_n.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",
    "            if m_name=='_rf_':\n",
    "                model = regressor(random_state=1, verbose=0)\n",
    "                clf = RandomizedSearchCV(\n",
    "                               model,\n",
    "                               param_grid_rf,\n",
    "                               verbose=0,\n",
    "                               n_iter=250,\n",
    "                               n_jobs=-1,\n",
    "                               cv=inner_cv.split(X_tr, y_tr),\n",
    "                              )\n",
    "\n",
    "                clf.fit(X_tr, y_tr.values.ravel())\n",
    "\n",
    "            else:\n",
    "                model = regressor(random_state=1, verbose=-1)\n",
    "                clf = RandomizedSearchCV(\n",
    "                               model,\n",
    "                               param_grid_lgbm,\n",
    "                               verbose=-1,\n",
    "                               n_iter=250,\n",
    "                               n_jobs=-1,\n",
    "                               cv=inner_cv.split(X_tr, y_tr),\n",
    "                              )\n",
    "                clf.fit(X_tr, y_tr, callbacks=None)\n",
    "\n",
    "\n",
    "            # predict using the best model\n",
    "            best_model = clf.best_estimator_\n",
    "            pred = best_model.predict(X_tt)\n",
    "            dff = pd.DataFrame({'Test':y_tt.values.squeeze(), 'Pred':pred}).reset_index(drop=True)\n",
    "            dff.to_csv(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/\"+model_var+\"_ensemble/CV_\"+str(j)+\"_\"+model_var+m_name+str(m)+\".csv\")\n",
    "\n",
    "            j+=1\n",
    "        #-----End of Nested CV ---------------------------------------------------\n",
    "        \n",
    "        if os.path.exists('/g/data/os22/chad_tmp/NEE_modelling/results/models_uncertainty/'+model_var+'/'+m_name+str(m)+'.joblib'):\n",
    "                print('    fit model and export')\n",
    "                continue\n",
    " \n",
    "        # Now conduct a hyperparameter test on all the data\n",
    "        # (minus the two removed sites) and fit a model\n",
    "        print('    fit model and export')\n",
    "        if m_name=='_rf_':\n",
    "            model = regressor(random_state=1, verbose=0)\n",
    "            clf = RandomizedSearchCV(\n",
    "                   model,\n",
    "                   param_grid_rf,\n",
    "                   verbose=0,\n",
    "                   n_iter=500,\n",
    "                   n_jobs=-1,\n",
    "                   cv=zip(train, test)\n",
    "                  )\n",
    "\n",
    "            clf.fit(x_n, y_n.values.ravel())    \n",
    "            model = regressor(**clf.best_params_)\n",
    "            model.fit(x_n, y_n.values.ravel())\n",
    "\n",
    "        else:\n",
    "            model = regressor(random_state=1, verbose=-1)\n",
    "            clf = RandomizedSearchCV(\n",
    "                   model,\n",
    "                   param_grid_lgbm,\n",
    "                   verbose=-1,\n",
    "                   n_iter=500,\n",
    "                   n_jobs=-1,\n",
    "                   cv=inner_cv.split(X_tr, y_tr),\n",
    "                  )\n",
    "            clf.fit(x_n, y_n, callbacks=None)\n",
    "            model = regressor(**clf.best_params_)\n",
    "            model.fit(x_n, y_n)    \n",
    "\n",
    "        dump(model, '/g/data/os22/chad_tmp/NEE_modelling/results/models_uncertainty/'+model_var+'/'+m_name+str(m)+'.joblib')\n",
    "\n",
    "    i+=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create 1:1 plots and statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_var='NEE'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base='/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/'+model_var+'_ensemble/'\n",
    "\n",
    "df_list=[]\n",
    "r2_list=[]\n",
    "ac_list=[]\n",
    "#get the list of cvs corresponding with a given CV split\n",
    "csvs = [i for i in os.listdir(base) if i.endswith('.csv')]\n",
    "for i in csvs:\n",
    "    df = pd.read_csv(base+i, usecols=['Test', 'Pred'])\n",
    "    obs,pred = df['Test'].values, df['Pred'].values\n",
    "    slope, intercept, r_value, p_value, std_err = stats.linregress(obs,pred)\n",
    "    r2_list.append(r_value**2)\n",
    "    ac_list.append(mean_absolute_error(obs, pred))\n",
    "    df_list.append(df)\n",
    "    \n",
    "\n",
    "#concantenate all the data  for the given CV split\n",
    "cross_df = pd.concat(df_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs,pred = cross_df['Test'].values, cross_df['Pred'].values\n",
    "slope, intercept, r_value, p_value, std_err = stats.linregress(obs,pred)\n",
    "r2 = r_value**2\n",
    "ac = mean_absolute_error(obs, pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,1, figsize=(6,6))\n",
    "font=19\n",
    "cross_df = cross_df.sample(n=10000)\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.3, 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",
    "if model_var != 'ET':\n",
    "    ax.set_xlabel('Observation '+ model_var + ' gC m\\N{SUPERSCRIPT TWO} mon⁻¹', fontsize=font)\n",
    "    ax.set_ylabel('Prediction ' + model_var+ ' gC m\\N{SUPERSCRIPT TWO} mon⁻¹', fontsize=font)\n",
    "\n",
    "else:\n",
    "    ax.set_xlabel('Observation '+ model_var + ' kgH\\N{SUBSCRIPT TWO}O/m\\N{SUPERSCRIPT TWO}/mon', fontsize=font)\n",
    "    ax.set_ylabel('Prediction ' + model_var+ ' kgH\\N{SUBSCRIPT TWO}O/m\\N{SUPERSCRIPT TWO}/mon', fontsize=font)\n",
    "    \n",
    "ax.text(.05, .95, 'r\\N{SUPERSCRIPT TWO}={:.2f}'.format(r2),\n",
    "            transform=ax.transAxes, fontsize=19)\n",
    "ax.text(.05, .9, 'MAE={:.3g}'.format(ac),\n",
    "            transform=ax.transAxes, fontsize=19)\n",
    "ax.tick_params(axis='x', labelsize=font)\n",
    "ax.tick_params(axis='y', labelsize=font)\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/cross_val_\"+model_var+\"_ensemble_\"+suffix+\".png\",\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ensemble feature importance\n",
    "\n",
    "Derived by calculating the mean absolute SHAP values for each feature in each model iteration, and subsequently averaging those values across all the models in the ensemble.\n",
    "\n",
    "RF feature importance takes a very long time!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import load\n",
    "import shap\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var='NEE'\n",
    "models_folder = '/g/data/os22/chad_tmp/NEE_modelling/results/models_uncertainty/'+var+'/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_list = [file for file in os.listdir(models_folder) if file.endswith(\".joblib\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#remove the columns we no longer need\n",
    "x = x.drop(['site'], axis=1)\n",
    "y = y.drop(['site'], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "dffs=[]\n",
    "i=1\n",
    "for m in model_list:\n",
    "    print(f\"{i}/{len(model_list)}\", end='\\r')\n",
    "    \n",
    "    if 'rf' in m:\n",
    "        #continue\n",
    "        explainer = shap.TreeExplainer(model = load(models_folder+m), data=x)\n",
    "    \n",
    "    else:\n",
    "        explainer = shap.Explainer(model = load(models_folder+m))\n",
    "    \n",
    "    shap_values = explainer(x)\n",
    "    vals = np.abs(shap_values.values).mean(0)\n",
    "    df = pd.DataFrame(list(zip(x.columns, vals)), columns=['col_name',m[0:-7]+'_FI'])\n",
    "    df.sort_values(by=[m[0:-7]+'_FI'],ascending=False,inplace=True)\n",
    "    df['col_name'] = df['col_name'].str.removesuffix(\"_RS\")\n",
    "    df = df.set_index('col_name', drop=True)\n",
    "    dffs.append(df)\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.concat(dffs, axis=1)\n",
    "df['mean'] = df.mean(axis=1)\n",
    "# df['stddev'] = df.std(axis=1)\n",
    "df = df.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(10,7))\n",
    "sb.barplot(data = df.sort_values(by='mean', axis=1, ascending=False).iloc[:, : 5].iloc[0:30], orient='h', ax=ax, palette='Blues_r')\n",
    "ax.tick_params(axis='x', labelsize=22)\n",
    "ax.tick_params(axis='y', labelsize=22)\n",
    "ax.set_xlabel('Ensemble avg. of mean abs. SHAP values: '+ var, fontsize=22)\n",
    "ax.set_ylabel('')\n",
    "plt.tight_layout()\n",
    "fig.savefig(\"/g/data/os22/chad_tmp/NEE_modelling/results/cross_validation/feature_importance_\"+var+\"_ensemble_\"+suffix+\".png\",\n",
    "            bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}