LSTM-ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import keras\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.stats import pearsonr\n",
    "from sklearn.preprocessing import MinMaxScaler, StandardScaler\n",
    "from keras.models import Sequential, Model\n",
    "from scipy.stats import pearsonr\n",
    "from keras.callbacks import EarlyStopping, ModelCheckpoint\n",
    "from keras.backend.tensorflow_backend import set_session\n",
    "from keras import backend as K\n",
    "from keras.backend import slice\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error\n",
    "from scipy.stats import pearsonr\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.model_selection import train_test_split\n",
    "import time\n",
    "import pickle\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(station): ## Station in String\n",
    "    flow = np.load('../usgsflow_'+station+'.npy')\n",
    "    precip = np.load('../NLDAS_precip_'+station+'.npy')\n",
    "    srad = np.load('../NLDAS_srad_'+station+'.npy')\n",
    "    tmax = np.load('../NLDAS_tmax_'+station+'.npy')\n",
    "    y = np.array(flow).reshape(-1, 1)\n",
    "    indx = np.where(y>=0)[0]\n",
    "    # print(precip.shape)\n",
    "    date = np.load('../usgsdate_'+station+'.npy', allow_pickle=True)\n",
    "    x = np.concatenate((precip, srad, tmax), axis=1)\n",
    "    x = x[indx]\n",
    "    y = y[indx]\n",
    "    date = date[indx]\n",
    "    return x, y\n",
    "def nse(y_pred, y_true):\n",
    "    nse = 1-np.sum((y_pred-y_true)**2)/np.sum((y_true-np.mean(y_true))**2)\n",
    "    return nse\n",
    "def dataset_ld(x,y,W,L):\n",
    "    obs = x.shape[0]\n",
    "    features = x.shape[1]\n",
    "    a = np.zeros([obs-W-L+1, W, features])\n",
    "    b = np.zeros([obs-W-L+1, 1])\n",
    "    for i in range(obs-W-L+1):\n",
    "        a[i,:,:] = x[i:i+W,:]\n",
    "        b[i,:] = y[i+W+L-1,0]    \n",
    "    return a, b\n",
    "def train_test_pre(x, y):\n",
    "    xtrain = x[:10000]; xtest = x[10000:]\n",
    "    ytrain = y[:10000]; ytest = y[10000:]\n",
    "    xscale = StandardScaler().fit(xtrain)\n",
    "    yscale = StandardScaler().fit(ytrain)\n",
    "    Xtrain = xscale.transform(xtrain); Xtest = xscale.transform(xtest)\n",
    "    Ytrain = yscale.transform(ytrain); Ytest = yscale.transform(ytest)\n",
    "    return Xtrain, Xtest, Ytrain, Ytest, xscale, yscale\n",
    "def custom_loss(y_true, y_pred):\n",
    "    s1 = K.sum((y_pred-y_true)**2)/K.sum((y_true-K.mean(y_true))**2)\n",
    "    return s1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(W,L):\n",
    "    model = Sequential()\n",
    "    model.add(keras.layers.CuDNNLSTM(50, return_sequences=True))\n",
    "    model.add(keras.layers.CuDNNLSTM(50, return_sequences=True))\n",
    "    model.add(keras.layers.CuDNNLSTM(50, return_sequences=False))\n",
    "    model.add(keras.layers.Dropout(0.5))\n",
    "    model.add(keras.layers.Dense(100))\n",
    "    model.add(keras.layers.Dense(1))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.0005; W=365; L=0;\n",
    "f = open('../../StationArea.pkl','rb')\n",
    "areas = pickle.load(f); f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# stations = np.load('../station-list.npy')\n",
    "stations = np.load('../first-stations.npy')\n",
    "d_nse = np.zeros((15,20)); d_mse = np.zeros((15,20)); d_mae = np.zeros((15,20)); d_r = np.zeros((15,20));\n",
    "high_mse = np.zeros((15,20))\n",
    "low_mse = np.zeros((15,20))\n",
    "i=0\n",
    "for station in stations:\n",
    "    x, y = load_data(str(station))\n",
    "    area = areas[str(station)]\n",
    "    ## Transform to Runoff\n",
    "    y = y*86400*1000/(area*1000*1000)\n",
    "    a_nse = []; a_mse = []; a_r = []; a_mae = []; total_time=0\n",
    "    a_high_mse=[]; a_low_mse=[]\n",
    "    best_nse = 0; model_name='LSTM/'+str(station)+'_LSTM.h5' ## Save the best nse and best model. \n",
    "    Xtrain, Xtest, Ytrain, Ytest, xscale, yscale = train_test_pre(x, y)\n",
    "    X_train, Y_train = dataset_ld(Xtrain, Ytrain, W, L)\n",
    "    X_test, Y_test = dataset_ld(Xtest, Ytest, W, L)\n",
    "    for training_id in range(15):\n",
    "        ensemble_name = 'LSTM/'+str(station)+'_LSTM_'+str(training_id)+'.h5'\n",
    "        model = build_model(W,L)\n",
    "        adam = keras.optimizers.Adam(lr=lr)\n",
    "        # model.compile(loss='mse', optimizer=adam)\n",
    "        model.compile(loss=custom_loss, optimizer=adam)\n",
    "        \n",
    "        # X_train, X_val, Y_train, Y_val = train_test_split(X_train, Y_train, test_size=0.25, random_state=7)\n",
    "        # training\n",
    "        start = time.time()\n",
    "        history = model.fit(X_train, Y_train, epochs=150, batch_size=512, \n",
    "                            verbose=0, shuffle=True)\n",
    "        run_time = time.time()-start\n",
    "        total_time+=run_time\n",
    "        # testing:\n",
    "        Y_pred = model.predict(X_test)\n",
    "        y_pred = yscale.inverse_transform((Y_pred).reshape(-1, 1))\n",
    "        y_true = yscale.inverse_transform((Y_test).reshape(-1, 1))\n",
    "        NSE = nse(y_pred, y_true); \n",
    "        R = pearsonr(y_pred.flatten(), y_true.flatten())[0]\n",
    "        MSE = mean_squared_error(y_pred, y_true); MAE = mean_absolute_error(y_pred, y_true) \n",
    "        a_nse+=[NSE]; a_mse+=[MSE]; a_r+=[R]; a_mae+=[MAE]\n",
    "        ## Save the best nse and best model. \n",
    "        if (NSE>best_nse):\n",
    "            best_nse=NSE; best_model=model; print('better')\n",
    "        model.save_weights(ensemble_name)\n",
    "        ## High flow and low flow\n",
    "        ind = np.argwhere(y_true<=np.percentile(y_true, 5))\n",
    "        low_pred = y_pred[ind]\n",
    "        low_true = y_true[ind]\n",
    "        ind = np.argwhere(y_true>=np.percentile(y_true, 95))\n",
    "        high_pred = y_pred[ind]\n",
    "        high_true = y_true[ind]\n",
    "        \n",
    "        e = high_pred-high_true; mse = np.mean(np.square(e)); a_high_mse+=[mse]\n",
    "        e = low_pred-low_true; mse = np.mean(np.square(e)); a_low_mse+=[mse]\n",
    "        \n",
    "        del model\n",
    "        del adam\n",
    "        \n",
    "    print(station,': run time is ', total_time/15, 's')\n",
    "    # print('NSE: ', a_nse, ' R: ', a_r)\n",
    "    # print('MSE: ', a_mse, ' MAE: ', a_mae)\n",
    "    d_nse[:,i]=a_nse; d_mse[:,i]=a_mse; d_mae[:,i]=a_mae; d_r[:,i]=a_r\n",
    "    high_mse[:,i]=a_high_mse; low_mse[:,i]=a_low_mse;\n",
    "    i+=1\n",
    "    # model.save(model_name)\n",
    "    best_model.save_weights(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('LSTM_NSE', d_nse);\n",
    "np.save('LSTM_MSE', d_mse);\n",
    "np.save('LSTM_high_mse', high_mse)\n",
    "np.save('LSTM_low_mse', low_mse)\n",
    "print('done')"
   ]
  }
 ],
 "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.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}