{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed hydrological modelling\n",
    "\n",
    "## Using Ravenpy to build a distributed hydrological model\n",
    "\n",
    "In this notebook, we will demonstrate how to build a distributed hydrological model using Raven as well as \"Routing product\" (Generated by BasinMaker), a database of subbasins and how they link to one another in a river network. Currently, Routing product is only available for North American catchments. However, if in time it becomes available on a larger scale, it would be trivial to change the setup apply it to other supported regions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "\n",
    "from ravenpy import Emulator\n",
    "from ravenpy.config import commands as rc\n",
    "from ravenpy.config import emulators\n",
    "from ravenpy.extractors.routing_product import (\n",
    "    BasinMakerExtractor,\n",
    "    open_shapefile,\n",
    "    upstream_from_coords,\n",
    ")\n",
    "\n",
    "# Utility that simplifies working with test data hosted on GitHub\n",
    "from ravenpy.testing.utils import get_file\n",
    "\n",
    "tmp_path = Path(tempfile.mkdtemp())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next step, we will get the Routing product file for our catchment. These can be downloaded here: http://hydrology.uwaterloo.ca/basinmaker/download_regional.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get path to a pre-downloaded BasinMaker Routing product database for our catchment\n",
    "shp_path = get_file(\"basinmaker/drainage_region_0175_v2-1/finalcat_info_v2-1.zip\")\n",
    "\n",
    "# Note that for this to work, the coordinates must be in the small\n",
    "# BasinMaker example (drainage_region_0175)\n",
    "df = open_shapefile(shp_path)\n",
    "\n",
    "# Gauge station for observations at Matapedia\n",
    "# SubId: 175000128\n",
    "# -67.12542 48.10417\n",
    "sub = upstream_from_coords(-67.12542, 48.10417, df)\n",
    "\n",
    "# Extract the subbasins and HRUs (one HRU per sub-basin)\n",
    "bm = BasinMakerExtractor(\n",
    "    df=sub,\n",
    "    hru_aspect_convention=\"ArcGIS\",\n",
    ")\n",
    "\n",
    "# Get the .rvh file that we will provide to the config and that links HRUs/subbasins to the river network\n",
    "rvh = bm.extract(hru_from_sb=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the HRUs and river network all setup, let's get the hydrometeorological data. We first get the database of streamflows and then do the same for weather. You can provide your own for your own catchments, here we are using our datasets to keep things tidy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Streamflow observations file\n",
    "qobs_fn = get_file(\"matapedia/Qobs_Matapedia_01BD009.nc\")\n",
    "\n",
    "# Make an observation gauge from the observed streamflow\n",
    "qobs = rc.ObservationData.from_nc(qobs_fn, alt_names=(\"discharge\",))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now prepare the meteorological data using the Gauge format. Note that this dataset of stations is a combination of stations that we iterate on, making a Gauge object for each station in our dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Meteo observations file\n",
    "meteo_grid_fn = get_file(\"matapedia/Matapedia_meteo_data_stations.nc\")\n",
    "\n",
    "# Alternate names for variables in the files\n",
    "alt_names = {\n",
    "    \"TEMP_MIN\": \"tmin\",\n",
    "    \"TEMP_MAX\": \"tmax\",\n",
    "    \"PRECIP\": \"pr\",\n",
    "}\n",
    "\n",
    "# Make virtual Gauges\n",
    "meteo_forcing_stations = [\n",
    "    rc.Gauge.from_nc(\n",
    "        meteo_grid_fn,\n",
    "        data_type=alt_names.keys(),\n",
    "        station_idx=i + 1,\n",
    "        alt_names=alt_names,\n",
    "    )\n",
    "    for i in range(6)  # Since we have 6 stations\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the data, we can run the distributed model as usual. Note that we must provide the AVG_ANNUAL_RUNOFF parameter to initialize the catchment's hydrological states for distributed models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the model configuration\n",
    "model_config = emulators.GR4JCN(\n",
    "    params=[0.529, -3.396, 407.29, 1.072, 16.9, 0.947],\n",
    "    StartDate=dt.datetime(1998, 1, 1),\n",
    "    EndDate=dt.datetime(2020, 12, 31),\n",
    "    ObservationData=[qobs],\n",
    "    Gauge=meteo_forcing_stations,\n",
    "    GlobalParameter={\"AVG_ANNUAL_RUNOFF\": 40.65},\n",
    "    **rvh,\n",
    ")\n",
    "\n",
    "# Run the model with the configuration we just built\n",
    "distributed_outputs = Emulator(model_config).run(overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explore the results, just like for any other model. However, this time we have a few gauges because the Routing Product integrates some gauges already. We want data for the first gauge:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the hydrograph object\n",
    "display(distributed_outputs.hydrograph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the resulting streamflow\n",
    "distributed_outputs.hydrograph.q_sim.isel(nbasins=0).plot.line(\n",
    "    x=\"time\", label=\"Distributed model\", color=\"blue\", lw=1.5\n",
    ")\n",
    "\n",
    "# Plot the observed streamflow\n",
    "qobs_data = xr.open_dataset(qobs_fn)\n",
    "qobs_data.discharge.plot.line(x=\"time\", label=\"Observations\", color=\"red\", lw=1.5)\n",
    "\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}