example/demo-LSTM-Tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM Tutorial for SMAP Soil Moisture\n",
    "\n",
    "This is a tutorial for training and testing a LSTM model for SMAP soil moisture prediction over CONUS."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Load needed packages**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'torch'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mModuleNotFoundError\u001B[0m                       Traceback (most recent call last)",
      "\u001B[1;32m<ipython-input-1-38fa65e14bfb>\u001B[0m in \u001B[0;36m<module>\u001B[1;34m\u001B[0m\n\u001B[0;32m      2\u001B[0m \u001B[0msys\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mpath\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mappend\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;34m'../'\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m      3\u001B[0m \u001B[1;32mimport\u001B[0m \u001B[0mos\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m----> 4\u001B[1;33m \u001B[1;32mimport\u001B[0m \u001B[0mtorch\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m      5\u001B[0m \u001B[1;32mimport\u001B[0m \u001B[0mnumpy\u001B[0m \u001B[1;32mas\u001B[0m \u001B[0mnp\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m      6\u001B[0m \u001B[1;32mfrom\u001B[0m \u001B[0mhydroDL\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mmaster\u001B[0m \u001B[1;32mimport\u001B[0m \u001B[0mdefault\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mwrapMaster\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mtrain\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;31mModuleNotFoundError\u001B[0m: No module named 'torch'"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from hydroDL.master import default, wrapMaster, train\n",
    "from hydroDL.data import dbCsv\n",
    "from hydroDL.post import plot, stat\n",
    "from hydroDL import master"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Define training options**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cDir = os.getcwd()\n",
    "# define training options\n",
    "optData = default.update(\n",
    "    default.optDataSMAP,\n",
    "    rootDB=os.path.join(cDir, 'data'),\n",
    "    subset='CONUSv4f1',\n",
    "    tRange=[20150401, 20160401])\n",
    "if torch.cuda.is_available():\n",
    "    optModel = default.optLstm\n",
    "else:\n",
    "    optModel = default.update(\n",
    "        default.optLstm,\n",
    "        name='hydroDL.model.rnn.CpuLstmModel')\n",
    "optLoss = default.optLossRMSE\n",
    "optTrain = default.update(\n",
    "    default.optTrainSMAP, \n",
    "    nEpoch=100,\n",
    "    saveEpoch=50)\n",
    "out = os.path.join(cDir, 'output', 'CONUSv4f1')\n",
    "masterDict = wrapMaster(out, optData, optModel, optLoss, optTrain)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Train the LSTM model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(masterDict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing part\n",
    "\n",
    "**This part is for model test and result plotting. By default the model will be saved in [here](output/CONUSv4f1/).**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - **Define test option**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = os.path.join(cDir, 'output', 'CONUSv4f1')\n",
    "rootDB = os.path.join(cDir, 'data')\n",
    "nEpoch = 100\n",
    "tRange = [20160401, 20170401]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - **Test the model in another year**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df, yp, yt = master.test(\n",
    "    out, tRange=[20160401, 20170401], subset='CONUSv4f1', epoch=nEpoch, reTest=True)\n",
    "yp = yp.squeeze()\n",
    "yt = yt.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - **Calculate statistic metrics.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate statistics\n",
    "statErr = stat.statError(yp, yt)\n",
    "# Box plots to show the test results\n",
    "statDictLst = [statErr]\n",
    "keyLst=['Bias', 'RMSE', 'ubRMSE', 'Corr']\n",
    "dataBox = list()\n",
    "for iS in range(len(keyLst)):\n",
    "    statStr = keyLst[iS]\n",
    "    temp = list()\n",
    "    for k in range(len(statDictLst)):\n",
    "        data = statDictLst[k][statStr]\n",
    "        data = data[~np.isnan(data)]\n",
    "        temp.append(data)\n",
    "    dataBox.append(temp)\n",
    "%matplotlib notebook\n",
    "fig = plot.plotBoxFig(dataBox, label1=keyLst, sharey=False, figsize=(12, 5))\n",
    "fig.patch.set_facecolor('white')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Plot an interactive map and users can click on map to show time series of observation and model predictions.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataGrid = [statErr['RMSE'], statErr['Corr']]\n",
    "dataTs = [yp, yt]\n",
    "t = df.getT()\n",
    "crd = df.getGeo()\n",
    "mapNameLst = ['RMSE', 'Correlation']\n",
    "tsNameLst = ['LSTM', 'SMAP']\n",
    "colorMap = None\n",
    "colorTs = None\n",
    "# plot map and time series\n",
    "%matplotlib notebook\n",
    "plot.plotTsMap(\n",
    "    dataGrid,\n",
    "    dataTs,\n",
    "    lat=crd[0],\n",
    "    lon=crd[1],\n",
    "    t=t,\n",
    "    mapNameLst=mapNameLst,\n",
    "    tsNameLst=tsNameLst,\n",
    "    isGrid=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}