overland_flow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab8e9d20",
   "metadata": {},
   "source": [
    "# Data Component Use Case for Overland Flow Simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec6d4e9",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Overland flow, particularly the infiltration-excess mechanism, is affected by the nature of water input through precipitation. This Jupyter notebook demonstrates how to use the [CSDMS Data Components](https://csdms.colorado.edu/wiki/DataComponents) to download the topography and precipitation datasets and use the [Landlab](https://landlab.readthedocs.io/en/master/) components to delineate the watershed and simulate the overland flow for a study area in the Boulder County. \n",
    "\n",
    "In this notebook, it includes the following sections:\n",
    "- [Initial Setup](#setup)\n",
    "  \n",
    "  Install API key files and create the input/output folders.\n",
    "  <br>\n",
    "- [Step 1: Download Dataset](#step1) \n",
    "\n",
    "  Download the topography and the precipitation datasets.\n",
    "  <br>\n",
    "- [Step 2: Delineate Watershed](#step2)\n",
    "\n",
    "  Use the topography dataset and Landlab components to delineate the watershed for the study area. \n",
    "  <br>\n",
    "- [Step 3: Calculate overland flow](#step3)\n",
    "\n",
    "  Use the watershed topography data and overland flow landlab component to simulate the surface water depth and discharge.\n",
    "  <br>\n",
    "- [Step 4: Visualize Results](#step4)\n",
    "\n",
    "  Visualize the final results as a short video.\n",
    "  <br>\n",
    "\n",
    "\n",
    "**Suggested Citation:** Gan, T., Tucker, G. E., Overeem, I. (2023). Data Component Use Case for Overland Flow Simulation, HydroShare, https://www.hydroshare.org/resource/bcbcfe823cc4432f8ce96c3048d4591f/\n",
    "\n",
    "**Run this notebook:** Please follow the instructions [here](https://github.com/gantian127/overlandflow_usecase#data-component-use-case-for-overland-flow-simulation) to run this notebook on the local PC or the online platform.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca52c800",
   "metadata": {},
   "source": [
    "<a id='setup'></a>\n",
    "## Initial Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2022db80",
   "metadata": {},
   "source": [
    "### Install API key files\n",
    "For the Topography and ERA5 data components, there is a need to create API key files to download the datasets. The install_api_key( ) function will ask for your [CDS API Key](https://cds.climate.copernicus.eu/api-how-to) and [Open Topography API Key](https://opentopography.org/blog/introducing-api-keys-access-opentopography-global-datasets) to create API key files. Please make sure you have already obtained those API Keys before you run this helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cef6fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import install_api_key\n",
    "install_api_key()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f582231",
   "metadata": {},
   "source": [
    "### Create folders\n",
    "We will first import all the python packages and then create three folders for this notebook:\n",
    "- **configuration file folder**: this folder includes a configuration file which will be used by the data component. In this example, we have prepared a configuration file ('dem_config.yaml') and put it in this folder. \n",
    "- **cache folder**: this folder stores the downloaded data file.\n",
    "- **results folder**: this folder stores the final results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b4ac8ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages \n",
    "import os\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray\n",
    "from datetime import datetime\n",
    "import cftime\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "import imageio.v2 as imageio\n",
    "from IPython.display import Video\n",
    "from tqdm import trange\n",
    "\n",
    "from pymt.models import Topography, Era5\n",
    "from landlab import RasterModelGrid, imshow_grid, imshow_grid_at_node\n",
    "from landlab.components import FlowAccumulator, ChannelProfiler\n",
    "from landlab.components.overland_flow import OverlandFlow\n",
    "from landlab.plot.drainage_plot import drainage_plot\n",
    "from landlab.utils import get_watershed_mask\n",
    "\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "plt.rcParams.update({'font.size': 14})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8a890a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create folders\n",
    "study_area = 'calwood'\n",
    "\n",
    "config_dir = os.path.join(os.getcwd(), 'config_files_{}'.format(study_area))\n",
    "results_dir = os.path.join(os.getcwd(), 'results_{}'.format(study_area)) \n",
    "cache_dir = os.path.join(os.getcwd(),'cache_{}'.format(study_area))\n",
    "\n",
    "\n",
    "for folder in [config_dir, results_dir, cache_dir]:\n",
    "    if not os.path.isdir(folder):\n",
    "        os.mkdir(folder)\n",
    "        print(folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a395cd3",
   "metadata": {},
   "source": [
    "<a id='step1'></a>\n",
    "## Step 1 Download Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2089ea7a",
   "metadata": {},
   "source": [
    "### OpenTopography DEM \n",
    "\n",
    "The number of wildfire increased in the Colorado State in recent years. A destructive wildfire can burn away plants and trees that prevent erosion. If heavy rains occur after such a fire, landslides and flash floods can occur. In this example, we will simulate the overland flow for a study area in the northwest of Boulder which was impacted by the CalWood Fire occurred in 2020. \n",
    "\n",
    "We will use the [Topography data component](https://csdms.colorado.edu/wiki/Model:Topography_Data_Component) to download the Digital Elevation Model (DEM) data with 30m resolution. The 'dem_config.yaml' file includes the parameter settings of this data component. The following cells demonstrate how to use the configuration file to initialize a data component and how to use the variable and grid related methods of this data component to get the metadata as well as the data values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03029206",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize Topography data component\n",
    "dem = Topography()\n",
    "dem.initialize(os.path.join(config_dir, 'dem_config.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aaf8732",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get DEM variable info\n",
    "var_name = dem.output_var_names[0]\n",
    "var_unit = dem.var_units(var_name)\n",
    "var_location = dem.var_location(var_name)\n",
    "var_type = dem.var_type(var_name)\n",
    "var_grid = dem.var_grid(var_name)\n",
    "var_itemsize = dem.var_itemsize(var_name)\n",
    "var_nbytes = dem.var_nbytes(var_name)\n",
    "print('variable_name: {} \\nvar_unit: {} \\nvar_location: {} \\nvar_type: {} \\nvar_grid: {} \\nvar_itemsize: {}' \n",
    "            '\\nvar_nbytes: {} \\n'. format(var_name, var_unit, var_location, var_type, var_grid, var_itemsize, var_nbytes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81bb67e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get DEM grid info \n",
    "dem_grid_ndim = dem.grid_ndim(var_grid) \n",
    "dem_grid_type = dem.grid_type(var_grid)\n",
    "dem_grid_shape = dem.grid_shape(var_grid)\n",
    "dem_grid_spacing = dem.grid_spacing(var_grid)\n",
    "dem_grid_origin = dem.grid_origin(var_grid)\n",
    "\n",
    "print('grid_ndim: {} \\ngrid_type: {} \\ngrid_shape: {} \\ngrid_spacing: {} \\ngrid_origin: {}'.format(\n",
    "    dem_grid_ndim, dem_grid_type, dem_grid_shape, dem_grid_spacing, dem_grid_origin))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a356fa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get DEM variable data\n",
    "dem_data = dem.get_value(var_name)\n",
    "dem_data_2D = dem_data.reshape(dem_grid_shape)\n",
    "\n",
    "# get X, Y extent for plot\n",
    "min_y, min_x = dem_grid_origin\n",
    "max_y = min_y + dem_grid_spacing[0]*(dem_grid_shape[0]-1)\n",
    "max_x = min_x + dem_grid_spacing[1]*(dem_grid_shape[1]-1)\n",
    "dy = dem_grid_spacing[0]/2\n",
    "dx = dem_grid_spacing[1]/2\n",
    "dem_extent = [min_x - dx, max_x + dx, min_y - dy, max_y + dy]\n",
    "\n",
    "# plot DEM data\n",
    "fig, ax = plt.subplots(1,1,figsize=(10,5))\n",
    "im = ax.imshow(dem_data_2D, extent=dem_extent)\n",
    "ax.title.set_text('Topography Data')\n",
    "ax.set_xlabel('latitude [degree_north]')\n",
    "ax.set_ylabel('longitude [degree_east]')\n",
    "fig.colorbar(im,label='elevation(m)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eaed0f03",
   "metadata": {},
   "source": [
    "### ERA5 Precipitation\n",
    "We will use the [ERA5 data component](https://csdms.colorado.edu/wiki/Model:ERA5_Data_Component) to download the hourly total precipitation data with 0.25 degrees (27-28km) resolution. The 'era5_config.yaml' file includes the parameter settings of this data component. The following cells demonstrate how to use the configuration file to initialize an ERA5 data component and how to use the variable, grid and time related methods to get the metadata as well as the data values. \n",
    "\n",
    "You'll notice that although the ERA5 and Topography data components download the datasets from different sources, they are using the same methods to get information from the datasets. Please note that sometimes the request for ERA5 data may be queued which may take a while (>10min) to get the data downloaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79708cbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize ERA5 data component\n",
    "era5 = Era5()\n",
    "era5.initialize(os.path.join(config_dir,'era5_config.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c9f128c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get ERA5 variable info\n",
    "var_name = era5.output_var_names[0]\n",
    "var_unit = era5.var_units(var_name)\n",
    "var_location = era5.var_location(var_name)\n",
    "var_type = era5.var_type(var_name)\n",
    "var_grid = era5.var_grid(var_name)\n",
    "var_itemsize = era5.var_itemsize(var_name)\n",
    "var_nbytes = era5.var_nbytes(var_name)\n",
    "print('variable_name: {} \\nvar_unit: {} \\nvar_location: {} \\nvar_type: {} \\nvar_grid: {} \\nvar_itemsize: {}' \n",
    "        '\\nvar_nbytes: {} \\n'. format(var_name, var_unit, var_location, var_type, var_grid, var_itemsize, var_nbytes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a7c38a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get ERA5 grid info\n",
    "era5_grid_ndim = era5.grid_ndim(var_grid) \n",
    "era5_grid_type = era5.grid_type(var_grid)\n",
    "era5_grid_shape = era5.grid_shape(var_grid)\n",
    "era5_grid_spacing = era5.grid_spacing(var_grid)\n",
    "era5_grid_origin = era5.grid_origin(var_grid)\n",
    "\n",
    "print('grid_ndim: {} \\ngrid_type: {} \\ngrid_shape: {} \\ngrid_spacing: {} \\ngrid_origin: {}'.format(\n",
    "    era5_grid_ndim, era5_grid_type, era5_grid_shape, era5_grid_spacing, era5_grid_origin))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93f4e644",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get ERA5 time info\n",
    "era5_start_time = era5.start_time\n",
    "era5_end_time = era5.end_time\n",
    "era5_time_step = era5.time_step\n",
    "era5_time_unit = era5.time_units\n",
    "era5_time_steps = int((era5_end_time - era5_start_time)/era5_time_step) + 1\n",
    "\n",
    "print('start_time:{} \\nend_time:{} \\ntime_step:{} \\ntime_unit:{} \\ntime_steps:{}'.format(\n",
    "    era5_start_time, era5_end_time, era5_time_step, era5_time_unit, era5_time_steps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f150435",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get ERA5 variables data (at first time step)\n",
    "era5_data = era5.get_value(var_name)\n",
    "era5_data_2D = era5_data.reshape(era5_grid_shape)\n",
    "\n",
    "# get X, Y extent for plot\n",
    "min_y, min_x = era5_grid_origin\n",
    "max_y = min_y + era5_grid_spacing[0]*(era5_grid_shape[0]-1)\n",
    "max_x = min_x + era5_grid_spacing[1]*(era5_grid_shape[1]-1)\n",
    "dy = era5_grid_spacing[0]/2\n",
    "dx = era5_grid_spacing[1]/2\n",
    "era5_extent = [min_x - dx, max_x + dx, min_y - dy, max_y + dy]\n",
    "\n",
    "# plot 2D ERA5 data\n",
    "fig, ax = plt.subplots(1,1,figsize=(10,5))\n",
    "im = ax.imshow(era5_data_2D, extent=era5_extent, cmap='Blues')\n",
    "ax.title.set_text('Precipitation Data'.format(var_unit))\n",
    "ax.set_xlabel('longitude [degree_east]')\n",
    "ax.set_ylabel('latitude [degree_north]')\n",
    "fig.colorbar(im,label='total precipitation(m/hr)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1641bfc7",
   "metadata": {},
   "source": [
    "The plot above shows the hourly total precipitation on June 25, 2021 (for the first time step at 00:00). The use case study site (see details in [Step2](#step2)) locates within the area that is covered by the pixel at the bottom left. \n",
    "\n",
    "We also create a time series plot to show the ERA5 hourly total precipitation data for June 25. In the plot below, we could see that the maximum hourly total precipitation is around 2.6 mm. This is much smaller than the observational rain intensity data obtained from the local station on June 25, which are 59.2 mm/hr (15min  intensity), 29.6 mm/hr (30min intensity), and 14.8 mm/hr (60min intensity). The value discrepancy is mainly caused by the fact that the ERA5 dataset has 0.25 degrees resolution and the total precipitation is calculated based on a large area leading to a very small value. Thus, to better represent the rain intensity situation at the study site, we will use the observational data from the local station for the overland flow simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee15f99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get ERA5 time series data\n",
    "time_steps = 24 \n",
    "prec_array = np.empty(time_steps)\n",
    "cftime_array = np.empty(time_steps)\n",
    "\n",
    "for i in range(0, time_steps):   \n",
    "    prec_array[i] = era5.get_value(var_name)[2] * 1000  # unit conversion m to mm \n",
    "    cftime_array[i] = era5.time\n",
    "    era5.update()\n",
    "\n",
    "time_array = cftime.num2date(cftime_array, era5_time_unit, only_use_cftime_datetimes=False, only_use_python_datetimes=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93838863",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot ERA5 time series\n",
    "plt.figure(figsize=(11,7))\n",
    "plt.plot(time_array, prec_array)\n",
    "plt.ylabel('total precipitation (mm)')\n",
    "plt.title('Time Series of Precipitation Data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6db88a",
   "metadata": {},
   "source": [
    "<a id='step2'></a>\n",
    "## Step 2 Delineate Watershed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883f3513",
   "metadata": {},
   "source": [
    "For our example, we are using the Heil Ranch area comprising Geer Canyon watershed nearby Boulder Colorado, USA. The watershed of Geer Canyon was affected in October 2020 by the Calwood wildfire. \n",
    "\n",
    "The Calwood Fire initiated approximate 3 miles North West of Jamestown, in Boulder County, Colorado. It started on a windy day, October 17th, 2020, and ran until contained on November 11th, 2020. The state of Colorado was in drought at the time of occurrence. High winds and dry conditions together caused the fire to spread rapidly, consuming nearly 1000 acres an hour.  This wildfire in total burned ~10,113 acres and destroyed 20 residential homes. It ranged in elevation between ~2500m of the highest terrain to ~1700m. It is the wildfire of the largest extent to occur in Boulder county, but affected less people and structures than the Marshall Fire in 2021.\n",
    "\n",
    "<img src=\"https://github.com/gantian127/overlandflow_usecase/blob/master/site.jpeg?raw=true\" width=\"900\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48612588",
   "metadata": {},
   "source": [
    "We will use the topography data for watershed delineation. The Landlab components ([FlowAccumulator](https://landlab.readthedocs.io/en/master/reference/components/flow_accum.html) and [ChannelProfiler](https://landlab.readthedocs.io/en/master/reference/components/channel_profiler.html)) and the Landlab utility ([get_watershed_mask](https://landlab.readthedocs.io/en/master/reference/utils/watershed.html)) will be used to accomplish this task. The general steps for watershed delineation include: \n",
    "- Setup raster model grid and add topographic elevation data field\n",
    "- Calculate the flow accumulation\n",
    "- Get watershed mask\n",
    "- Set watershed boundary "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9025c6cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up raster model grid  \n",
    "model_grid = RasterModelGrid(dem_grid_shape, xy_spacing=30)\n",
    "\n",
    "# add topographic elevation data field\n",
    "dem_field = model_grid.add_field(\"topographic__elevation\", dem_data.astype('float'))\n",
    "\n",
    "# plot the topographic elevation data field\n",
    "plt.figure(figsize=(10,5))\n",
    "ax = plt.gca()\n",
    "ax.invert_yaxis()\n",
    "imshow_grid(model_grid,\n",
    "            'topographic__elevation',\n",
    "            plot_name='Data field of topographic elevation',\n",
    "            var_name='elevation(m)')\n",
    "ax.set_xlabel('east-west distance (m)')\n",
    "ax.set_ylabel('north-south distance (m)')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0753fa37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the flow accumulation\n",
    "fa=FlowAccumulator( model_grid, method='Steepest', \n",
    "                    flow_director='FlowDirectorSteepest', \n",
    "                    depression_finder='LakeMapperBarnes',\n",
    "                    redirect_flow_steepest_descent=True,\n",
    "                    reaccumulate_flow=True)\n",
    "fa.run_one_step()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07577c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the flow accumulation result\n",
    "plt.figure(figsize=(10,5))\n",
    "ax = plt.gca()\n",
    "ax.invert_yaxis()\n",
    "imshow_grid(\n",
    "    model_grid,\n",
    "    np.sqrt(model_grid.at_node[\"drainage_area\"]),\n",
    "    plot_name=\"Data field of drainage area\",\n",
    "    var_name=\"square root of drainage area\",\n",
    ")\n",
    "ax.set_xlabel('east-west distance (m)')\n",
    "ax.set_ylabel('north-south distance (m)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b3c330c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up channel profiler \n",
    "profiler = ChannelProfiler(model_grid,number_of_watersheds=1)\n",
    "profiler.run_one_step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a16f12b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get watershed mask\n",
    "outlet = profiler.nodes[0][0]\n",
    "watershed_mask = get_watershed_mask(model_grid, outlet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b639cd15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set watershed boundary\n",
    "model_grid.at_node[\"topographic__elevation\"][~watershed_mask] = -9999.0\n",
    "model_grid.status_at_node[~watershed_mask] = model_grid.BC_NODE_IS_CLOSED\n",
    "model_grid.status_at_node[outlet] = model_grid.BC_NODE_IS_FIXED_VALUE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fdd266b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the watershed \n",
    "plt.figure(figsize=(10,5))\n",
    "ax = plt.gca()\n",
    "ax.invert_yaxis()\n",
    "imshow_grid_at_node(model_grid, dem_field, \n",
    "                    plot_name='Topographic elevation of the watershed',\n",
    "                    colorbar_label=\"elevation (m)\")\n",
    "\n",
    "ax.set_xlabel('east-west distance (m)')\n",
    "ax.set_ylabel('north-south distance (m)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe2f469",
   "metadata": {},
   "source": [
    "<a id='step3'></a>\n",
    "## Step 3  Calculate Overland Flow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d67cf94",
   "metadata": {},
   "source": [
    "In this step, the topography elevation of the watershed and the [Overland Flow component](https://landlab.readthedocs.io/en/master/user_guide/overland_flow_user_guide.html#background-on-overlandflow-component) will be used for the simulation. The model run time is set as 200 min with the first 10 min assigned with a constant rainfall intensity (59.2 mm/hr), which is based on the observation on June 25, 2021 (the summer after the Calwood fire occured). \n",
    "\n",
    "In this simulation, we suppose there is no water infiltration process involved. Is this justified? Not really, this is a model simplification for demonstration purposes. But, wildfires cause soils to become extremely dry and  hydrophobic, and decrease infiltration. So this assumption provides an end-member case of extreme soil water repellency affecting the Geer Canyon watershed.\n",
    "\n",
    "The simulation results will create a discharge time series plot at the outlet and 2D plots of the surface water depth at each 30 sec time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f708f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add surface wate depth data field \n",
    "model_grid.add_zeros(\"surface_water__depth\", at=\"node\", clobber=True)\n",
    "\n",
    "# set initial surface water depth value\n",
    "model_grid.at_node[\"surface_water__depth\"].fill(\n",
    "    1e-12\n",
    ") \n",
    "\n",
    "plt.figure(figsize=(10,5))\n",
    "ax = plt.gca()\n",
    "ax.invert_yaxis()\n",
    "imshow_grid(model_grid,'surface_water__depth', \n",
    "            plot_name='Data field of surface water depth',\n",
    "            var_name='water depth (m)',\n",
    "            cmap=\"Blues\", vmin=0.0,vmax=1)\n",
    "ax.set_xlabel('east-west distance (m)')\n",
    "ax.set_ylabel('north-south distance (m)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f28c0e9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate overland flow component\n",
    "overland_flow = OverlandFlow(\n",
    "    model_grid, steep_slopes=True\n",
    ")\n",
    "\n",
    "# set model run parameters\n",
    "elapsed_time = 0.0\n",
    "model_run_time = 200 * 60  # duration of run (s) \n",
    "storm_duration = 10 * 60  # duration of rain (s)  \n",
    "time_step = 120 \n",
    "rainfall_intensity = 59.2 /(1000*3600)  # mm/hr to m/s\n",
    "\n",
    "outlet_discharge = []\n",
    "outlet_times = []\n",
    "\n",
    "# run overland flow simulation\n",
    "for time_slice in trange(time_step, model_run_time + time_step, time_step):\n",
    "\n",
    "    while elapsed_time < time_slice:\n",
    "\n",
    "        # get adaptive time step\n",
    "        overland_flow.dt = min(overland_flow.calc_time_step(), time_step)\n",
    "\n",
    "        # set rainfall intensity\n",
    "        if elapsed_time < (storm_duration):\n",
    "            overland_flow.rainfall_intensity = rainfall_intensity\n",
    "        else:\n",
    "            overland_flow.rainfall_intensity = 0.0\n",
    "\n",
    "        # run model\n",
    "        overland_flow.overland_flow(dt=overland_flow.dt)\n",
    "\n",
    "        # update elapsed time\n",
    "        elapsed_time += overland_flow.dt\n",
    "\n",
    "        # get discharge result\n",
    "        discharge = overland_flow.discharge_mapper(\n",
    "            model_grid.at_link[\"surface_water__discharge\"], convert_to_volume=True\n",
    "        )\n",
    "\n",
    "        outlet_discharge.append(discharge[outlet])\n",
    "        outlet_times.append(elapsed_time)\n",
    "        \n",
    "    # plot result\n",
    "    fig, ax = plt.subplots(\n",
    "        2, 1, figsize=(8, 9), gridspec_kw={\"height_ratios\": [1, 1.5]}\n",
    "    )\n",
    "    fig.suptitle(\"Results at {} min\".format(time_slice / 60))\n",
    "\n",
    "    ax[0].plot(outlet_times, outlet_discharge, \"-\")\n",
    "    ax[0].set_xlabel(\"Time elapsed (s)\")\n",
    "    ax[0].set_ylabel(\"discharge (cms)\")\n",
    "    ax[0].set_title(\"Water discharge at the outlet\")\n",
    "\n",
    "    imshow_grid(\n",
    "        model_grid,\n",
    "        \"surface_water__depth\",\n",
    "        cmap=\"Blues\",\n",
    "        vmin=0,\n",
    "        vmax=1.2,\n",
    "        var_name=\"surface water depth (m)\",\n",
    "    )\n",
    "    ax[1].set_title(\"\")\n",
    "    ax[1].invert_yaxis()\n",
    "    ax[1].set_xlabel('east-west distance (m)')\n",
    "    ax[1].set_ylabel('north-south distance (m)')\n",
    "\n",
    "    plt.close(fig)\n",
    "    fig.savefig(os.path.join(results_dir, \"flow_{}.png\".format(time_slice)))\n",
    "    \n",
    "print('Simulation is done')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e20afbe",
   "metadata": {},
   "source": [
    "<a id='step4'></a>\n",
    "## Step 4 Visualize Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2142cf54",
   "metadata": {},
   "source": [
    "Run the cells below and it will show a short video. You will be able to see the changes of the surface water depth over the watershed and the discharge at the outlet over the simulation time. \n",
    "\n",
    "You can also reset the value of the \"rain_intensity\" variable using the ERA5 dataset (see the 2nd code cell in [Step3](#step3)). Then rerun Step3 and Step4 to create new results for comparison. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "123aef54",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Make a short video\n",
    "img_files = [os.path.join(results_dir, file) for file in os.listdir(results_dir) if '.png' in file]\n",
    "img_files.sort(key=lambda x: os.path.getmtime(x))\n",
    "\n",
    "with imageio.get_writer(os.path.join(results_dir,'overland_flow.mp4'), mode='I', fps=1,macro_block_size=None ) as writer:\n",
    "    for f in img_files:\n",
    "        im = imageio.imread(os.path.join(results_dir, f))\n",
    "        writer.append_data(im)\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a67bfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the video\n",
    "Video(\"./results_calwood/overland_flow.mp4\", width=900, height=700)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd64bbc",
   "metadata": {},
   "source": [
    "## References\n",
    "- Adams, J. M., Gasparini, N. M., Hobley, D. E. J., Tucker, G. E., Hutton, E. W. H., Nudurupati, S. S., and Istanbulluoglu, E. (2017). The Landlab v1.0 OverlandFlow component: a Python tool for computing shallow-water flow across watersheds, Geoscientific Model Development, 10, 1645–1663, https://doi.org/10.5194/gmd-10-1645-2017.\n",
    "\n",
    "- Barnhart, K. R., Hutton, E. W. H., Tucker, G. E., M. Gasparini, N., Istanbulluoglu, E., E. J. Hobley, D., J. Lyons, N., Mouchene, M., Siddhartha Nudurupati, S., M. Adams, J., & Bandaragoda, C. (2020). Short communication: Landlab v2.0: A software package for Earth surface dynamics. Earth Surface Dynamics, 8(2), 379–397. https://doi.org/10.5194/esurf-8-379-2020\n",
    "\n",
    "- de Almeida, G. A., Bates, P., Freer, J. E., & Souvignet, M. (2012). Improving the stability of a simple formulation of the shallow water equations for 2‐D flood modeling. Water Resources Research, 48(5). https://doi.org/10.1029/2011WR011570"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ebbd082",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}