{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluating Ensemble Streamflow Forecasts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook introduces using TEEHR to evaluate ensemble forecasts. Note that ensemble analysis in TEEHR is currently ongoing, and additional metrics and analyses are in development.\n",
    "\n",
    "In this exercise we will use a pre-existing evaluation stored in the TEEHR repository consisting of two USGS gage locations with [Hydrologic Ensemble Forecasting System (HEFS)](https://www.weather.gov/abrfc/about_HEFS) forecasts from [CNRFC.](https://www.cnrfc.noaa.gov/ensembleHourlyProductCSV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read evaluation data from S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "import teehr\n",
    "from teehr.examples.setup_ensemble_example import setup_hefs_example\n",
    "\n",
    "# Tell Bokeh to output plots in the notebook\n",
    "from bokeh.io import output_notebook\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the directory where the Evaluation will be created.\n",
    "test_eval_dir = Path(Path().home(), \"temp\", \"09_ensemble_metrics\")\n",
    "shutil.rmtree(test_eval_dir, ignore_errors=True)\n",
    "\n",
    "# Setup the example evaluation using data from the TEEHR repository.\n",
    "setup_hefs_example(tmpdir=test_eval_dir)\n",
    "\n",
    "# Initialize the evaluation.\n",
    "ev = teehr.Evaluation(dir_path=test_eval_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore the evaluation data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we have two USGS gage locations corresponding to CNRFC forecast points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "locations_gdf = ev.locations.to_geopandas()\n",
    "locations_gdf.teehr.locations_map()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The secondary location IDs correspond to the CNRFC forecast point IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev.location_crosswalks.to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The primary timeseries are the USGS gage observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_df = ev. primary_timeseries.to_pandas()\n",
    "pt_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_df.teehr.timeseries_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The secondary timeseries are the CNRFC HEFS ensemble forecasts corresponding to each location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st_df = ev.secondary_timeseries.to_pandas()\n",
    "st_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HEFS is probabilistic and each forecast consists of several ensemble traces or members. In TEEHR,each trace is represented by the `member` column in the timeseries table schema."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case each forecast contains 42 members."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st_df.member.unique().size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the data. We'll use some helper functions to plot the forecasts against the USGS gage data. Note that visualization functionality within TEEHR is in-development."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import hvplot.pandas  # noqa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_teehr_timeseries(\n",
    "    df: pd.DataFrame,\n",
    "    x_column: str = \"value_time\",\n",
    "    y_column: str = \"value\",\n",
    "    title: str = None\n",
    "):\n",
    "    columns = df.columns\n",
    "    variable_columns = [col for col in columns if df[col].nunique() > 1]\n",
    "    variable_columns.remove(x_column)\n",
    "    variable_columns.remove(y_column)\n",
    "    figure = df.hvplot(x=\"value_time\", y=\"value\", by=variable_columns)\n",
    "    return figure.options(\n",
    "        width=700,\n",
    "        height=400,\n",
    "        show_grid=True,\n",
    "        title=title\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the HEFS forecast for location SCBC1.\n",
    "st_df = ev.secondary_timeseries.query(\"location_id == 'SCBC1'\").to_pandas()\n",
    "\n",
    "# Create a line plot of the corresponding USGS observed flow data.\n",
    "usgs_plot = pt_df[pt_df.location_id == \"usgs-11402000\"].hvplot(x=\"value_time\", y=\"value\").options(line_width=3, line_dash=\"dashed\", color=\"black\")\n",
    "\n",
    "# Create a line plot of the HEFS forecast data with the USGS observed flow data.\n",
    "(plot_teehr_timeseries(st_df, title=\"SCBC1 HEFS Ensemble vs. USGS\") * usgs_plot).options(show_legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the HEFS forecast for location MRYC1.\n",
    "st_df = ev.secondary_timeseries.query(\"location_id == 'MRYC1'\").to_pandas()\n",
    "\n",
    "# Create a line plot of the corresponding USGS observed flow data.\n",
    "usgs_plot = pt_df[pt_df.location_id == \"usgs-11421000\"].hvplot(x=\"value_time\", y=\"value\").options(line_width=3, line_dash=\"dashed\", color=\"black\")\n",
    "\n",
    "# Create a line plot of the HEFS forecast data with the USGS observed flow data.\n",
    "(plot_teehr_timeseries(st_df, title=\"MRYC1 HEFS Ensemble vs. USGS\") * usgs_plot).options(show_legend=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating the Continuous Ranked Probability Score (CRPS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The continuous ranked probability score (CRPS) is a common metric used to evaluate the\n",
    "performance of ensemble forecasts. TEEHR implements CRPS through the [scoringrules](https://frazane.github.io/scoringrules/) library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define a CRPS metric object and customize it's parameters in the same way we defined metrics in previous examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "crps = teehr.ProbabilisticMetrics.CRPS()\n",
    "\n",
    "# Specify the estimator to use.\n",
    "#   \"pwm\" - Probability Weighted Moment form (default)\n",
    "#   \"nrg\" - Energy form\n",
    "#   \"fair\" - Fair version of energy form\n",
    "crps.estimator = \"pwm\"\n",
    "\n",
    "# Specify the backend to use.\n",
    "#   \"numba\" - Use the numba library for calculations (default)\n",
    "#   \"numpy\" - Use the numpy library for calculations\n",
    "crps.backend = \"numba\"\n",
    "\n",
    "crps.output_field_name = \"crps_ensemble\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev.joined_timeseries.create()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev.metrics.query(\n",
    "    include_metrics=[crps],\n",
    "    group_by=[\n",
    "        \"primary_location_id\",\n",
    "        \"reference_time\",\n",
    "        \"configuration_name\"\n",
    "    ],\n",
    "    order_by=[\"primary_location_id\"],\n",
    ").to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, CRPS returns an array of scores corresponding to each value time. We can summarize these values by defining a function (typically mean) in the metric object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "crps.summary_func = np.mean\n",
    "\n",
    "crps.output_field_name = \"mean_crps_ensemble\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev.metrics.query(\n",
    "    include_metrics=[crps],\n",
    "    group_by=[\n",
    "        \"primary_location_id\",\n",
    "        \"reference_time\",\n",
    "        \"configuration_name\"\n",
    "    ],\n",
    "    order_by=[\"primary_location_id\"],\n",
    ").to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently only the CRPS is implemented in TEEHR however additonal metrics and functionality is ongoing including ensemble timeseries visualization, skill score calculation, Brier Score, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "ev.spark.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}