ExploreFLPE_EU.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/SteveCoss/SWOTdawgDISTRO/blob/main/ExploreFLPE_EU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "44e87c43",
      "metadata": {
        "id": "44e87c43"
      },
      "source": [
        "# Notebook to explore flow law parameter estimation \n",
        "\n",
        "Flow law parameter estimation (FLPE, pronounced \"floppy\") is the heart of Confluence. There are two stages. Stage 1 algorithms function at reach scale: this has been much of the heart of the DAWG's work over the past years. They function at the scale of a reach or several reaches. Stage 2 is the so-called integrator, which functions at the scale of basins or continents. \n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os, sys\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "nb_path = '/content/notebooks'\n",
        "os.symlink('/content/drive/My Drive/DAWGnotebooks/Path_files', nb_path)\n",
        "sys.path.insert(0,nb_path)"
      ],
      "metadata": {
        "id": "XVfLmWH9E7B2"
      },
      "id": "XVfLmWH9E7B2",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#reset working directory to distro folder\n",
        "!pwd\n",
        "import os\n",
        "os.chdir(\"/content/drive/My Drive/DAWGnotebooks/dist_5.0\")\n",
        "!pwd"
      ],
      "metadata": {
        "id": "aJgKuezmE70Q"
      },
      "id": "aJgKuezmE70Q",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "319b2f2d",
      "metadata": {
        "tags": [],
        "id": "319b2f2d"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from netCDF4 import Dataset\n",
        "from pathlib import Path\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import json\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "591ddd0b",
      "metadata": {
        "id": "591ddd0b"
      },
      "source": [
        "This notebook relies on the swotdawgviz library. Use either the version included in this package or use an up-to-date version (https://github.com/klarnier/swotdawgviz)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "90f06c2e",
      "metadata": {
        "id": "90f06c2e"
      },
      "outputs": [],
      "source": [
        "# Using embedded version of swotdawgviz\n",
        "from swotdawgviz.swotdawgviz import io as sdvio\n",
        "from swotdawgviz.swotdawgviz import maps as sdvm\n",
        "from swotdawgviz.swotdawgviz import plots as sdvpl\n",
        "\n",
        "# Using up-to-date installed version of swotdawgviz\n",
        "# from swotdawgviz import io as sdvio\n",
        "# from swotdawgviz import maps as sdvm\n",
        "# from swotdawgviz import plots as sdvpl"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cae0d900",
      "metadata": {
        "id": "cae0d900"
      },
      "source": [
        "# Stage 1 FLPE\n",
        "\n",
        "There are six stage 1 FLPEs: BAM, HiVDI, MetroMan, MOMMA, SAD and SIC4DVar. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4e6e44d8",
      "metadata": {
        "id": "4e6e44d8"
      },
      "outputs": [],
      "source": [
        "# read in the list of reaches\n",
        "\n",
        "# specify input directory\n",
        "InputDir=Path('.')\n",
        "\n",
        "# read in file with all reaches in the domain\n",
        "reach_json=InputDir.joinpath('reaches.json')\n",
        "with open(reach_json) as json_file:\n",
        "    reaches = json.load(json_file)\n",
        "\n",
        "#extract reach ids\n",
        "domain_reachids=list()\n",
        "for reach in reaches:\n",
        "    domain_reachids.append(reach['reach_id'])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ae9546c3",
      "metadata": {
        "id": "ae9546c3"
      },
      "source": [
        "## GeoBAM\n",
        "\n",
        "GeoBAM applies mass conservation flow law inversion on all the nodes within a reach, and thus operates on one reach at a time, and all output is within one file at a time. The GeoBAM flow law is:\n",
        "\n",
        "$ \\log{W} = \\left( \\dfrac{3}{5}\\dfrac{1}{r} +1\\right) \\left( \\log{Q} +\\log{n} - \\dfrac{5}{3} \\log{D_b} + \\dfrac{5}{3} \\log{\\dfrac{r+1}{r}} + \\dfrac{5}{3} r \\log{W_b} - \\dfrac{1}{2} \\log{S} \\right)$\n",
        "\n",
        "where $W$ is the measured width, $S$ is the measured slope, $D_b$ and $W_b$ are the bankfull width and depth, respectively, $r$ is a channel shape parameter, and $n$ is the resistance coefficient."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cf0dcaf2",
      "metadata": {
        "id": "cf0dcaf2"
      },
      "outputs": [],
      "source": [
        "s1flpedir=InputDir.joinpath('s1-flpe')\n",
        "gbdir=s1flpedir.joinpath('geobam').joinpath('output')\n",
        "\n",
        "gb_con_dir=gbdir.joinpath('output_constrained')\n",
        "#gb_unc_dir=gbdir.joinpath('output_unconstrained')\n",
        "\n",
        "gbfiles=os.listdir(gb_con_dir)\n",
        "\n",
        "print(gbfiles[0:10]) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9f4153e6",
      "metadata": {
        "id": "9f4153e6"
      },
      "outputs": [],
      "source": [
        "# read and check out geobam data\n",
        "\n",
        "reach='21303000061'\n",
        "\n",
        "gbfile=gb_con_dir.joinpath(reach + '_geobam.nc') # this reach is on the Kentucky River\n",
        "\n",
        "gb_dataset = Dataset(gbfile)\n",
        "\n",
        "print(gb_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5928aa22",
      "metadata": {
        "id": "5928aa22"
      },
      "outputs": [],
      "source": [
        "# GeoBAM runs three different Markov Chains to estimate discharge. let's extract the estimated bankfull depth\n",
        "# nt=gb_dataset['nt'][:]\n",
        "# t=list(nt.data)\n",
        "# nt=len(t)\n",
        "\n",
        "bamDbs=np.empty( (3,)  )\n",
        "\n",
        "for c in range(3):\n",
        "    bamDbs[c]=gb_dataset['logDb']['mean' + str(c+1)][:].filled(np.nan)\n",
        "    \n",
        "bamDb = np.mean(bamDbs)\n",
        "\n",
        "gb_dataset.close()\n",
        "\n",
        "print('GeoBAM bankfull depth for ',reach,'=',np.exp(bamDb))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c968d5ab",
      "metadata": {
        "id": "c968d5ab"
      },
      "outputs": [],
      "source": [
        "# loop over reaches and extract depth\n",
        "Db_gb=dict()\n",
        "Db_gbs=list()\n",
        "\n",
        "for reach in domain_reachids:\n",
        "    gbfile=gb_con_dir.joinpath(str(reach) + '_geobam.nc') # this reach is on the Kentucky River\n",
        "    \n",
        "    Db_gb[str(reach)]=np.empty((1))\n",
        "    \n",
        "    try: \n",
        "        gb_dataset = Dataset(gbfile)\n",
        "    except: \n",
        "        Db_gb[str(reach)]=-1.    \n",
        "\n",
        "    if not Db_gb[str(reach)] == -1:                    \n",
        "        try:\n",
        "            data=gb_dataset['logDb']        \n",
        "        except:\n",
        "            Db_gb[str(reach)]=-1.    \n",
        "\n",
        "        if not Db_gb[str(reach)]==-1:        \n",
        "            for c in range(3):\n",
        "                Db_gb[str(reach)]+=np.exp(gb_dataset['logDb']['mean' + str(c+1)][:].filled(np.nan))/3        \n",
        "        \n",
        "        gb_dataset.close()\n",
        "        \n",
        "    Db_gbs.append(float(Db_gb[str(reach)]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "122f2683",
      "metadata": {
        "id": "122f2683"
      },
      "outputs": [],
      "source": [
        "# summary of geobam bankfull depth estimates\n",
        "fig,ax=plt.subplots()\n",
        "plt.hist(Db_gbs)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "be883007",
      "metadata": {
        "id": "be883007"
      },
      "outputs": [],
      "source": [
        "Db_gbs = np.array(Db_gbs)\n",
        "Db_gbs = Db_gbs[Db_gbs>-1]\n",
        "fig,ax=plt.subplots()\n",
        "plt.hist(Db_gbs)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "86b876b9",
      "metadata": {
        "id": "86b876b9"
      },
      "outputs": [],
      "source": [
        "# plot map of bankfull depth - first create a map to attach them to\n",
        "\n",
        "swot_nc_dir=InputDir.joinpath('swot/timeseries')\n",
        "\n",
        "# sword_shp_dir='/Users/mtd/Data/SWOT/SWORD/SWORD/Reaches_Nodes/shp/NA/'\n",
        "# sword_shp_dir='/home/kevin/Documents/SWOT/DATABASES/SWORD/V11/shp/NA/'\n",
        "\n",
        "sword_shp_dir=InputDir.joinpath('sword/shp/EU')\n",
        "\n",
        "\n",
        "collection = sdvio.SwotObservationsCollection(swot_nc_dir)\n",
        "\n",
        "sword_hb23_reaches = sdvio.SwordShapefile(os.path.join(sword_shp_dir, \"eu_sword_reaches_hb23_v11.shp\"),\n",
        "                                          reaches_list=collection.reaches_list)\n",
        "\n",
        "rmap = sdvm.ReachesMap(sword_hb23_reaches.dataset)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f5d6db1b",
      "metadata": {
        "id": "f5d6db1b"
      },
      "outputs": [],
      "source": [
        "# add bankfull depth the rmap object\n",
        "rmap._dataset['geobam_db']=-1.\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    if not np.isnan(Db_gb[str(reachid)]):\n",
        "        rmap._dataset.loc[rmap._dataset['reach_id'].astype(str)==str(reachid),['geobam_db']]=Db_gb[str(reachid)]\n",
        "    \n",
        "\n",
        "rmap._json_dataset = rmap._dataset.to_json()    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "847c0e3c",
      "metadata": {
        "id": "847c0e3c"
      },
      "outputs": [],
      "source": [
        "#map Geobam bankfull depth\n",
        "geobamdb_map = rmap.get_centerlines_map(varname=\"geobam_db\",varlimits=[0,2])\n",
        "geobamdb_map"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ed6a89b7",
      "metadata": {
        "id": "ed6a89b7"
      },
      "source": [
        "## MetroMan\n",
        "\n",
        "MetroMan runs on a so-called \"Inversion Set\" of reaches all at once, and so data is stored in files that contain output from all sets. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "04680a1e",
      "metadata": {
        "id": "04680a1e"
      },
      "outputs": [],
      "source": [
        "mmdir=s1flpedir.joinpath('metroman').joinpath('output').joinpath('output_constrained')\n",
        "mmfiles=os.listdir(mmdir)\n",
        "\n",
        "print(mmfiles[0:10]) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8741726",
      "metadata": {
        "id": "b8741726"
      },
      "outputs": [],
      "source": [
        "mmfile=mmdir.joinpath('24298000081-24298000071-24298000061_metroman.nc') \n",
        "\n",
        "mm_dataset = Dataset(mmfile)\n",
        "\n",
        "print(mm_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "60b12475",
      "metadata": {
        "id": "60b12475"
      },
      "outputs": [],
      "source": [
        "A0=mm_dataset['A0hat'][:].filled(np.nan)\n",
        "\n",
        "mm_dataset.close()\n",
        "\n",
        "fig,ax=plt.subplots()\n",
        "ax.plot(A0)\n",
        "plt.ylabel('Cross-sectional area at first overpass, m^2')\n",
        "plt.xlabel('Reach #')\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5d617472",
      "metadata": {
        "id": "5d617472"
      },
      "source": [
        "## MOMMA\n",
        "\n",
        "MOMMA files are one per reach. Flow law parameters are chosen to make the mean discharge match the prior, rather than solving the McFLI problem."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "be98a130",
      "metadata": {
        "id": "be98a130"
      },
      "outputs": [],
      "source": [
        "modir=s1flpedir.joinpath('momma').joinpath('output').joinpath('output_unconstrained')\n",
        "mofiles=os.listdir(modir)\n",
        "print(mofiles[0:10]) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "03be4baf",
      "metadata": {
        "id": "03be4baf"
      },
      "outputs": [],
      "source": [
        "# let's loop over these and compute a momma average dishcarge for each reach\n",
        "\n",
        "Qbar_momma=dict()\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    mofile=str(reachid) + '_momma.nc'\n",
        "    mofilepath=modir.joinpath(mofile)\n",
        "    \n",
        "    if mofilepath.exists():\n",
        "        mo_dataset = Dataset(mofilepath)\n",
        "        Qfile=mo_dataset['Q'][:].filled(np.nan)\n",
        "        Qbar_momma[str(reachid)]=np.nanmean(Qfile)\n",
        "        mo_dataset.close()        \n",
        "    else:\n",
        "        Qbar_momma[str(reachid)]=np.nan\n",
        "        "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d3cb6200",
      "metadata": {
        "id": "d3cb6200"
      },
      "outputs": [],
      "source": [
        "# create a map to show the MOMMA average discharge\n",
        "\n",
        "# add Momma discharge to the rmap object\n",
        "rmap._dataset['QMomma']=-1.\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    if not np.isnan(Qbar_momma[str(reachid)]):\n",
        "        rmap._dataset.loc[rmap._dataset['reach_id'].astype(str)==str(reachid),['QMomma']]=Qbar_momma[str(reachid)]\n",
        "    \n",
        "\n",
        "rmap._json_dataset = rmap._dataset.to_json()    \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c997485c",
      "metadata": {
        "tags": [],
        "id": "c997485c"
      },
      "outputs": [],
      "source": [
        "#map Momma Discharge\n",
        "momma_map = rmap.get_centerlines_map(varname=\"QMomma\",varlimits=[0,2000])\n",
        "momma_map"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a596e77a",
      "metadata": {
        "id": "a596e77a"
      },
      "source": [
        "## H2iVDI\n",
        "\n",
        "H2iVDI files are one per reach. Flow law parameters are computed to match the discharge infered by the Full St-Venant model using 4DVAR.\n",
        "\n",
        "<span style=\"color:red\"> Warning: H2iVDI only ran on 75/213 reaches (reaches with valid observations)</span>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "19569f77",
      "metadata": {
        "id": "19569f77"
      },
      "outputs": [],
      "source": [
        "h2ivdidir=s1flpedir.joinpath('hivdi').joinpath('output').joinpath('output_unconstrained')\n",
        "h2ivdifiles=os.listdir(h2ivdidir)\n",
        "print(h2ivdifiles[0:10]) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a447fc34",
      "metadata": {
        "id": "a447fc34"
      },
      "outputs": [],
      "source": [
        "# let's loop over these and retrieve H2iVDI unobserved flow area (A0)\n",
        "\n",
        "A0_h2ivdi=dict()\n",
        "\n",
        "count = 0\n",
        "for reachid in domain_reachids:\n",
        "    h2ivdifile=str(reachid) + '_hivdi.nc'\n",
        "    h2ivdifilepath=h2ivdidir.joinpath(h2ivdifile)\n",
        "    \n",
        "    if h2ivdifilepath.exists():\n",
        "        h2ivdi_dataset = Dataset(h2ivdifilepath)\n",
        "        reach_dataset = h2ivdi_dataset.groups[\"reach\"]\n",
        "        A0file=reach_dataset['A0'][0]\n",
        "        A0_h2ivdi[str(reachid)]=float(A0file)\n",
        "        h2ivdi_dataset.close()\n",
        "        count += 1\n",
        "    else:\n",
        "        A0_h2ivdi[str(reachid)]=np.nan\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f6933bee",
      "metadata": {
        "id": "f6933bee"
      },
      "outputs": [],
      "source": [
        "# create a map to show the H2iVDI unobserved flow area (A0)\n",
        "\n",
        "# add H2iVDI A0 to the rmap object\n",
        "rmap._dataset['A0H2iVDI']=-1.\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    if not np.isnan(A0_h2ivdi[str(reachid)]):\n",
        "        rmap._dataset.loc[rmap._dataset['reach_id'].astype(str)==str(reachid),['A0H2iVDI']]=A0_h2ivdi[str(reachid)]\n",
        "    \n",
        "\n",
        "rmap._json_dataset = rmap._dataset.to_json()    \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f7270e03",
      "metadata": {
        "tags": [],
        "id": "f7270e03"
      },
      "outputs": [],
      "source": [
        "#map H2iVDI unobserved flow area (A0)\n",
        "h2ivdi_map = rmap.get_centerlines_map(varname=\"A0H2iVDI\",varlimits=[1000,10000])\n",
        "h2ivdi_map"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3c85ee3b",
      "metadata": {
        "id": "3c85ee3b"
      },
      "source": [
        "Here we compare three products for H2iVDI. The first two products are the **constrained** and **unconstrained** products using the *classic* HiVDI toolchain (Full St-Venant+4DVAR) and the third is the new *surrogate* (hybrid) toolchain which uses machine learning technics to estimate better priors of discharge and $A_0$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "60b99723",
      "metadata": {
        "id": "60b99723"
      },
      "outputs": [],
      "source": [
        "#--------------------------------\n",
        "# REQUIRES UP-TO-DATE swotdawgviz\n",
        "#--------------------------------\n",
        "\n",
        "reachid = 74259000031\n",
        "# Load SoS data\n",
        "sosdir=InputDir.joinpath('sos')\n",
        "sosdir_constrained=sosdir.joinpath(\"constrained\")\n",
        "sosdir_unconstrained=sosdir.joinpath(\"unconstrained\")\n",
        "sos_unconstrained = sdvio.SosNetCDF(sosdir_unconstrained.joinpath(\"na_sword_v11_SOS_priors.nc\"), verbose=False)\n",
        "sos_constrained = sdvio.SosNetCDF(sosdir_constrained.joinpath(\"na_sword_v11_SOS_priors.nc\"), verbose=False)\n",
        "\n",
        "# Retrieve in-situ data and priors\n",
        "reach_data = sos_unconstrained.get_reach(reachid)\n",
        "wbm_mean_q = float(reach_data[\"model_mean_q\"])\n",
        "reach_data = sos_constrained.get_reach(reachid)\n",
        "usgs_mean_q = float(reach_data[\"usgs_mean_q\"])\n",
        "grades_mean_q = float(reach_data[\"model_mean_q\"])\n",
        "\n",
        "# Load HiVDI output\n",
        "h2ivdidir_constrained=s1flpedir.joinpath('hivdi').joinpath('output').joinpath('output_constrained')\n",
        "h2ivdi_constrained = sdvio.OutputH2iVDI(h2ivdidir_constrained.joinpath(\"%i_hivdi.nc\" % reachid))\n",
        "h2ivdidir_unconstrained=s1flpedir.joinpath('hivdi').joinpath('output').joinpath('output_unconstrained')\n",
        "h2ivdi_unconstrained = sdvio.OutputH2iVDI(h2ivdidir_unconstrained.joinpath(\"%i_hivdi.nc\" % reachid))\n",
        "h2ivdidir_surrogate=s1flpedir.joinpath('hivdi').joinpath('output').joinpath('surrogate')\n",
        "h2ivdi_surrogate = sdvio.OutputH2iVDI(h2ivdidir_surrogate.joinpath(\"%i_hivdi.nc\" % reachid))\n",
        "\n",
        "# Make plot\n",
        "statuses = \"const.:%i,unconst.:%i\" % (h2ivdi_constrained.status(), h2ivdi_unconstrained.status())\n",
        "statuses += \",surrogate:%i\" % h2ivdi_surrogate.status()\n",
        "qplot = sdvpl.DischargePlot(title='Reach %i, status:%s' % (reachid, statuses), date_units=\"days_since_2000\")\n",
        "qplot.add_prior(\"WBM mean Q\", wbm_mean_q, color=\"red\", linestyle=\"dash\")\n",
        "qplot.add_prior(\"GRADES mean Q\", grades_mean_q, color=\"red\", linestyle=None)\n",
        "if not np.isnan(usgs_mean_q):\n",
        "    qplot.add_prior(\"USGS mean Q\", usgs_mean_q, color=\"green\", linestyle=None)\n",
        "qplot.add_product(\"H2iVDI (constrained)\", h2ivdi_constrained.Q, times=h2ivdi_constrained.t,\n",
        "                  color=\"blue\", linestyle=None)\n",
        "qplot.add_product(\"H2iVDI (unconstrained)\", h2ivdi_unconstrained.Q, times=h2ivdi_unconstrained.t,\n",
        "                  color=\"blue\", linestyle=\"dash\")\n",
        "qplot.add_product(\"H2iVDI (surrogate)\", h2ivdi_surrogate.Q, times=h2ivdi_surrogate.t,\n",
        "                  color=\"cyan\", linestyle=\"dash\")\n",
        "qplot.render(backend=\"plotly\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "89d7b940",
      "metadata": {
        "tags": [],
        "id": "89d7b940"
      },
      "source": [
        "## Comparison of discharge estimated by 3 different algorithms\n",
        "Here we compare the discharge timeseries infered by algorithms MetroMan, MOMMA and H2iVDI on a single reach. Note that these are reading discharge created internally in each stage 1 FLPE algorithm in Confluence. These values will differ from those created after the integrator, and after the offline discharge is created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ff61af99",
      "metadata": {
        "id": "ff61af99"
      },
      "outputs": [],
      "source": [
        "reachid = 74264100171"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "52305b92",
      "metadata": {
        "id": "52305b92"
      },
      "outputs": [],
      "source": [
        "# Retrieve MetroMan discharge\n",
        "mmfile=mmdir.joinpath('74264100171-74264100161-74264100151_metroman.nc') # this set is on the Cumberland River\n",
        "mm_dataset = Dataset(mmfile)\n",
        "QMetroMan = mm_dataset[\"allq\"][0,:]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "daffbff5",
      "metadata": {
        "id": "daffbff5"
      },
      "outputs": [],
      "source": [
        "# Retrieve MOMMA discharge\n",
        "mofile=str(reachid) + '_momma.nc'\n",
        "mofilepath=modir.joinpath(mofile)\n",
        "if mofilepath.exists():\n",
        "    mo_dataset = Dataset(mofilepath)\n",
        "    QMomma = mo_dataset['Q'][:].filled(np.nan)\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "34fecb8c",
      "metadata": {
        "id": "34fecb8c"
      },
      "outputs": [],
      "source": [
        "# Retrieve H2iVDI discharge\n",
        "h2ivdifile=str(reachid) + '_hivdi.nc'\n",
        "h2ivdifilepath=h2ivdidir.joinpath(h2ivdifile)\n",
        "if h2ivdifilepath.exists():\n",
        "    h2ivdi_dataset = Dataset(h2ivdifilepath)\n",
        "    reach_dataset = h2ivdi_dataset.groups[\"reach\"]\n",
        "    QH2iVDI = reach_dataset['Q'][:].filled(np.nan)\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0c4357d7",
      "metadata": {
        "id": "0c4357d7"
      },
      "outputs": [],
      "source": [
        "# Retrieve GeoBam discharge\n",
        "gbfile=str(reachid) + '_geobam.nc'\n",
        "gbfilepath=gb_unc_dir.joinpath(gbfile)\n",
        "if gbfilepath.exists():\n",
        "    gb_dataset = Dataset(gbfilepath)    \n",
        "    q_dataset=gb_dataset.groups['q']\n",
        "    q1 = q_dataset['q1'][:].filled(np.nan)\n",
        "    q2 = q_dataset['q2'][:].filled(np.nan)\n",
        "    q3 = q_dataset['q3'][:].filled(np.nan)\n",
        "    qs=np.vstack((q1,q2,q3))\n",
        "    QGeoBAM=np.nanmean(qs, axis=0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "504e9451",
      "metadata": {
        "tags": [],
        "id": "504e9451"
      },
      "outputs": [],
      "source": [
        "# Compare the discharge timeseries\n",
        "fig,ax=plt.subplots()\n",
        "ax.plot(h2ivdi_dataset['nt'][:], QMetroMan, label=\"MetroMan\")\n",
        "ax.plot(h2ivdi_dataset['nt'][:], QMomma, label=\"MOMMA\")\n",
        "ax.plot(h2ivdi_dataset['nt'][:], QH2iVDI, label=\"H2iVDI\")\n",
        "ax.plot(h2ivdi_dataset['nt'][:], QGeoBAM, label=\"GeoBAM\")\n",
        "\n",
        "ax.set_ylim(0.0, 2500)\n",
        "\n",
        "plt.ylabel('Discharge, cms')\n",
        "plt.xlabel('Days since Jan. 2000')\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bc5626b6",
      "metadata": {
        "id": "bc5626b6"
      },
      "outputs": [],
      "source": [
        "h2ivdi_dataset['nt']"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c024a5c0",
      "metadata": {
        "id": "c024a5c0"
      },
      "source": [
        "# Stage 2 FLPE\n",
        "\n",
        "In this stage, we impose a couple of constraints on our retrievals. First, we require that discharge be conserved across river networks. Second, we use gages to constrain the discharge values. We accomplish this using the (still unpublished) algorithm called \"mean optimization integrator\" or MOI. MOI \"integrates\" information from across the river network to refine the flow law parameter estimates for both the constrained and the unconstrained branches of discharge."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "93401abd",
      "metadata": {
        "id": "93401abd"
      },
      "outputs": [],
      "source": [
        "# all discharge for all algorithms are stored in reach-by-reach integrator files\n",
        "\n",
        "stage2_dir=InputDir.joinpath('s2-flpe').joinpath('output')\n",
        "s2_files=os.listdir(stage2_dir)\n",
        "print(s2_files[0:10]) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ce35df91",
      "metadata": {
        "id": "ce35df91"
      },
      "outputs": [],
      "source": [
        "Qbar_mm2=dict()\n",
        "sbQrel=dict()\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    mm2_file=str(reachid) + '_integrator.nc'\n",
        "    mm2_filepath=stage2_dir.joinpath(mm2_file)\n",
        "    \n",
        "    if mm2_filepath.exists():\n",
        "        mm2_dataset = Dataset(mm2_filepath)\n",
        "        Qfile=mm2_dataset['metroman/qbar_basinScale'][:].filled(np.nan)\n",
        "        Qbar_mm2[str(reachid)]=np.nanmean(Qfile)\n",
        "        \n",
        "        sbQrel[str(reachid)]=mm2_dataset['metroman/sbQ_rel'][:].filled(np.nan)\n",
        "        \n",
        "        mm2_dataset.close()\n",
        "    else:\n",
        "        Qbar_mm2[str(reachid)]=np.nan\n",
        "        "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "de3b8a3a",
      "metadata": {
        "id": "de3b8a3a"
      },
      "outputs": [],
      "source": [
        "# add integrated MetroMan discharge and uncertainty to the rmap object\n",
        "rmap._dataset['QMetroMan2']=-1.\n",
        "rmap._dataset['sbQrel']=0.\n",
        "\n",
        "for reachid in domain_reachids:\n",
        "    if not np.isnan(Qbar_mm2[str(reachid)]):\n",
        "        rmap._dataset.loc[rmap._dataset['reach_id'].astype(str)==str(reachid),['QMetroMan2']]=Qbar_mm2[str(reachid)]\n",
        "        rmap._dataset.loc[rmap._dataset['reach_id'].astype(str)==str(reachid),['sbQrel']]=sbQrel[str(reachid)]\n",
        "    \n",
        "\n",
        "rmap._json_dataset = rmap._dataset.to_json()    \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "183526f8",
      "metadata": {
        "id": "183526f8"
      },
      "outputs": [],
      "source": [
        "# map MetroMan discharge\n",
        "mm2_map = rmap.get_centerlines_map(varname=\"QMetroMan2\",varlimits=[0,7500])\n",
        "mm2_map"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "633b07c5",
      "metadata": {
        "id": "633b07c5"
      },
      "outputs": [],
      "source": [
        "# map MetroMan discharge uncertainty\n",
        "mm_sbq_map = rmap.get_centerlines_map(varname=\"sbQrel\",varlimits=[0,.5])\n",
        "mm_sbq_map\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d06d4480",
      "metadata": {
        "id": "d06d4480"
      },
      "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.8.13"
    },
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "ed6a89b7",
        "5d617472",
        "a596e77a",
        "89d7b940"
      ],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}