{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A coupled rainfall-runoff model in Landlab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial demonstrates a very simple synthetic rainfall-runoff model in Landlab, using the `SpatialPrecipitationDistribution` and `OverlandFlow` components. This assumes no infiltration, but it could be added by modifying the `rainfall__flux` field appropriately.\n",
    "\n",
    "First, import the modules we'll need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import os\n",
    "import numpy as np\n",
    "from landlab.io import read_esri_ascii, write_esri_ascii\n",
    "from landlab import imshow_grid_at_node\n",
    "from landlab.components import SpatialPrecipitationDistribution\n",
    "from landlab.components import OverlandFlow\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up a grid and load some arbitrary existing catchment elevation data. A functional version of this might use a real gauged catchment for comparison to reality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# here we use an arbitrary, very small, \"real\" catchment\n",
    "fname = 'hugo_site.asc'\n",
    "mg, z = read_esri_ascii(fname, name='topographic__elevation')\n",
    "mg.status_at_node[mg.nodes_at_right_edge] = mg.BC_NODE_IS_FIXED_VALUE\n",
    "mg.status_at_node[np.isclose(z, -9999.)] = mg.BC_NODE_IS_CLOSED\n",
    "\n",
    "plt.figure()\n",
    "imshow_grid_at_node(mg, z, colorbar_label='Elevation (m)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build a mocked-up rainfall distribution using the `SpatialPrecipitationDistribution` component.\n",
    "\n",
    "It would be trivial to replace this with an imported real rainfall field - and we save and reload the pattern to highlight how this might work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rain = SpatialPrecipitationDistribution(mg)\n",
    "np.random.seed(26)  # arbitrary to get a cool-looking storm out every time\n",
    "\n",
    "# get the storm simulator to provide a storm\n",
    "# There's only one storm generated here in the time series, so easy enough to do.\n",
    "# first, check the directory we need for saving exists, and make it if not:\n",
    "if not os.path.exists('./rainfall'):\n",
    "    os.makedirs('./rainfall')\n",
    "for (storm_t, interstorm_t) in rain.yield_storms(style='monsoonal'):  # storm lengths in hrs\n",
    "    mg.at_node['rainfall__flux'] *= 0.001  # because the rainfall comes out in mm/h\n",
    "    mg.at_node['rainfall__flux'] *= 10.  # to make the storm heavier and more interesting!\n",
    "    plt.figure()\n",
    "    # plot up this storm\n",
    "    imshow_grid_at_node(\n",
    "        mg, 'rainfall__flux', cmap='gist_ncar', colorbar_label='Rainfall flux (m/h)'\n",
    "    )\n",
    "    plt.show()\n",
    "    write_esri_ascii('./rainfall/rainfall.asc', mg, 'rainfall__flux', clobber=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, load the rainfall files and set up the model, telling the flood router to accept the rainfalls in the file(s) as inputs. \n",
    "\n",
    "In the first instance, this is set up as an instantaneous storm, with all the water dropped over the catchment in one go. Below, we modify this assumption to allow time distributed rainfall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for filename in os.listdir('./rainfall'):  # for each file in the folder\n",
    "    if filename.endswith(\".asc\"):  # ...that ends with .asc...\n",
    "        # remove any rainfall field that already exists on the grid:\n",
    "        try:\n",
    "            _ = mg.at_node.pop('rainfall__flux')\n",
    "        except KeyError:\n",
    "            pass\n",
    "        _, q_rain = read_esri_ascii(\n",
    "            './rainfall/'+filename, grid=mg, name='rainfall__flux')\n",
    "    else:\n",
    "        continue\n",
    "    mg.add_zeros(\"surface_water__depth\", at=\"node\")   \n",
    "    mg.at_node['surface_water__depth'].fill(1.e-12)  # a veneer of water stabilises the model\n",
    "    mg.at_node['surface_water__depth'] += mg.at_node['rainfall__flux'] * storm_t\n",
    "    of = OverlandFlow(mg, steep_slopes=True)\n",
    "\n",
    "\n",
    "    # storm_t here is the duration of the rainfall, from the rainfall component\n",
    "    # We're going to assume the rainfall arrives effectively instantaneously, but\n",
    "    # adding discharge during the run is completely viable\n",
    "\n",
    "    node_of_max_q = 2126  # established by examining the output of a previous run\n",
    "    outlet_depth = []\n",
    "    outlet_times = []\n",
    "    post_storm_elapsed_time = 0.\n",
    "    last_storm_loop_tracker = 0.\n",
    "    while post_storm_elapsed_time < 0.5 * 3600.:  # plot 30 mins-worth of runoff\n",
    "        dt = of.calc_time_step()\n",
    "        of.run_one_step(dt=dt)\n",
    "        post_storm_elapsed_time += dt\n",
    "        storm_loop_tracker = post_storm_elapsed_time % 180.  # show every 3 min\n",
    "        # NB: Do NOT allow this plotting if there are multiple files in the folder\n",
    "        if storm_loop_tracker < last_storm_loop_tracker:\n",
    "            plt.figure()\n",
    "            imshow_grid_at_node(\n",
    "                mg,\n",
    "                'surface_water__depth',\n",
    "                var_name='Stage (m)')\n",
    "            plt.title('Stage at t=' + str(post_storm_elapsed_time//1) + 's')\n",
    "            plt.show()\n",
    "        last_storm_loop_tracker = storm_loop_tracker\n",
    "        outlet_depth.append(mg.at_node['surface_water__depth'][node_of_max_q])\n",
    "        outlet_times.append(post_storm_elapsed_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, plot the time series at the outlet (defined as the node that experiences peak stage):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(outlet_times, outlet_depth, '-')\n",
    "plt.xlabel('Time elapsed (s)')\n",
    "plt.ylabel('Flood stage (m)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can relax the assumption that all this discharge is delivered instantaneously at the start of the run with some tweaking of the driver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for filename in os.listdir('./rainfall'):  # for each file in the folder\n",
    "    if filename.endswith(\".asc\"):  # ...that ends with .asc...\n",
    "        # remove any rainfall field that already exists on the grid:\n",
    "        try:\n",
    "            _ = mg.at_node.pop('rainfall__flux')\n",
    "        except KeyError:\n",
    "            pass\n",
    "        _, q_rain = read_esri_ascii(\n",
    "            './rainfall/'+filename, grid=mg, name='rainfall__flux')\n",
    "    else:\n",
    "        continue\n",
    "      \n",
    "    mg.at_node['surface_water__depth'].fill(1.e-12)\n",
    "\n",
    "    of = OverlandFlow(mg, steep_slopes=True)\n",
    "    node_of_max_q = 2126\n",
    "    total_mins_to_plot = 60.  # plot 60 mins-worth of runoff\n",
    "    plot_interval_mins = 10.  # show every 10 min\n",
    "    min_tstep_val = 1.  # necessary to get the model going cleanly\n",
    "    outlet_depth = []\n",
    "    outlet_times = []\n",
    "    storm_elapsed_time = 0.\n",
    "    total_elapsed_time = 0.\n",
    "    last_storm_loop_tracker = 0.\n",
    "    while total_elapsed_time < total_mins_to_plot * 60.:\n",
    "        dt = of.calc_time_step()\n",
    "        remaining_total_time = total_mins_to_plot * 60. - total_elapsed_time\n",
    "        if storm_elapsed_time < storm_t * 3600.:\n",
    "            remaining_storm_time = storm_t * 3600. - storm_elapsed_time\n",
    "            dt = min((dt, remaining_total_time, remaining_storm_time, min_tstep_val))\n",
    "        else:\n",
    "            dt = min((dt, remaining_total_time, min_tstep_val))\n",
    "        of.run_one_step(dt=dt)\n",
    "        total_elapsed_time += dt\n",
    "        storm_elapsed_time += dt\n",
    "        storm_loop_tracker = total_elapsed_time % (plot_interval_mins * 60.)\n",
    "        # NB: Do NOT allow this plotting if there are multiple files in the folder\n",
    "        if storm_loop_tracker < last_storm_loop_tracker:\n",
    "            plt.figure()\n",
    "            imshow_grid_at_node(\n",
    "                mg,\n",
    "                'surface_water__depth',\n",
    "                var_name='Stage (m)')\n",
    "            plt.title('Stage at t=' + str(total_elapsed_time//1) + 's')\n",
    "            plt.show()\n",
    "        last_storm_loop_tracker = storm_loop_tracker\n",
    "        outlet_depth.append(mg.at_node['surface_water__depth'][node_of_max_q])\n",
    "        outlet_times.append(total_elapsed_time)\n",
    "        if storm_elapsed_time < storm_t * 3600.:\n",
    "            mg.at_node['surface_water__depth'] += mg.at_node['rainfall__flux'] * dt / 3600."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(outlet_times, outlet_depth, '-')\n",
    "plt.xlabel('Time elapsed (s)')\n",
    "plt.ylabel('Flood stage (m)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, a more realistic spread of the rainfall across the storm gives a longer and more subdued flood pulse.\n",
    "\n",
    "(An aside: the levelling off of the tail at h~0.125m is likely due to the permanent filling of a depression in the topography - the same thing is probably causing the deep pixels in the flow maps - or are these numerical instabilities? Resolving this is left as an exercise for the reader...)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}