{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Local frequency analyses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic imports\n",
    "import hvplot.xarray  # noqa\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xdatasets\n",
    "\n",
    "import xhydro as xh\n",
    "import xhydro.frequency_analysis as xhfa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data extraction and preparation\n",
    "\n",
    "In this example, we will perform a frequency analysis using historical timeseries data collected from various sites. Our first step is to acquire a hydrological dataset. For this, we utilize the [xdatasets](https://hydrologie.github.io/xdatasets/notebooks/getting_started.html) library to retrieve hydrological data from the [Ministère de l'Environnement, de la Lutte contre les changements climatiques, de la Faune et des Parcs](https://www.cehq.gouv.qc.ca/atlas-hydroclimatique/stations-hydrometriques/index.htm) in Québec, Canada. Our query will focus on stations with IDs beginning with `020`, and specifically those with natural flow pattern. Some stations have information on water levels, but will will only process streamflow data.\n",
    "\n",
    "Alternatively, users can provide their own `xarray.DataArray`. When preparing the data for frequency analysis, it must adhere to the following requirements:\n",
    "\n",
    "- The dataset must include a `time` dimension.\n",
    "- If there is a 1D spatial dimension (e.g., `id` in the example below), it must contain an attribute `cf_role` with the value `timeseries_id`.\n",
    "- The variable must have at least a `units` attribute. Although additional attributes, such as `long_name` and `cell_methods`, are not strictly required, they are expected by `xclim`, which is invoked during the frequency analysis. Missing these attributes will trigger warnings.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = (\n",
    "    xdatasets.Query(\n",
    "        **{\n",
    "            \"datasets\": {\n",
    "                \"deh\": {\n",
    "                    \"id\": [\"020*\"],\n",
    "                    \"regulated\": [\"Natural\"],\n",
    "                    \"variables\": [\"streamflow\"],\n",
    "                }\n",
    "            },\n",
    "            \"time\": {\"start\": \"1970-01-01\", \"minimum_duration\": (15 * 365, \"d\")},\n",
    "        }\n",
    "    )\n",
    "    .data.squeeze()\n",
    "    .load()\n",
    ")\n",
    "\n",
    "# This dataset lacks some of the aforementioned attributes, so we need to add them.\n",
    "ds = ds.rename({\"streamflow\": \"q\"})\n",
    "ds[\"id\"].attrs[\"cf_role\"] = \"timeseries_id\"\n",
    "ds[\"q\"].attrs = {\n",
    "    \"long_name\": \"Streamflow\",\n",
    "    \"units\": \"m3 s-1\",\n",
    "    \"standard_name\": \"water_volume_transport_in_river_channel\",\n",
    "    \"cell_methods\": \"time: mean\",\n",
    "}\n",
    "\n",
    "# Clean some of the coordinates that are not needed for this example\n",
    "ds = ds.drop_vars([c for c in ds.coords if c not in [\"id\", \"time\", \"name\"]])\n",
    "\n",
    "# For the purpose of this example, we keep only 2 stations\n",
    "ds = ds.isel(id=slice(3, 5))\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.q.dropna(\"time\", how=\"all\").hvplot(\n",
    "    x=\"time\", grid=True, widget_location=\"bottom\", groupby=\"id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Acquiring block maxima\n",
    "\n",
    "The `xhydro.indicators.get_yearly_op` function can be used to extract block maxima from a time series. This function provides several options for customizing the extraction process, such as selecting the desired time periods and more. In the following section, we will present the main arguments available to help tailor the block maxima extraction to your needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xh.indicators.get_yearly_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a) Defining seasons\n",
    "\n",
    "Seasons can be defined using indexers compatible with `xclim.core.calendar.select_time`. Four types of indexers are currently supported:\n",
    "\n",
    "- `month`: A sequence of month numbers (e.g., `[1, 2, 12]` for January, February, and December).\n",
    "- `season`: A sequence of season abbreviations, with options being `'DJF'`, `'MAM'`, `'JJA'`, and `'SON'` (representing Winter, Spring, Summer, and Fall, respectively).\n",
    "- `doy_bounds`: A sequence specifying the inclusive bounds of the period in terms of day of year (e.g., `[152, 243]`).\n",
    "- `date_bounds`: Similar to `doy_bounds`, but using a month-day format (`'%m-%d'`), such as `[\"01-15\", \"02-23\"]`.\n",
    "\n",
    "When using `xhydro.indicators.get_yearly_op` to calculate block maxima, the indexers should be grouped within a dictionary and passed to the `timeargs` argument. The dictionary keys should represent the requested period (e.g., `'winter'`, `'summer'`) and will be appended to the variable name. Each dictionary entry can include the following:\n",
    "\n",
    "- The indexer, as defined above (e.g., `\"date_bounds\": [\"02-11\", \"06-19\"]`).\n",
    "- (Optional) An annual resampling frequency. This is mainly used for indexers that span across the year. For example, setting `\"freq\": \"YS-DEC\"` will start the year in December instead of January."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some examples\n",
    "timeargs = {\n",
    "    \"spring\": {\"date_bounds\": [\"02-11\", \"06-19\"]},\n",
    "    \"summer\": {\"doy_bounds\": [152, 243]},\n",
    "    \"fall\": {\"month\": [9, 10, 11]},\n",
    "    \"winter\": {\n",
    "        \"season\": [\"DJF\"],\n",
    "        \"freq\": \"YS-DEC\",\n",
    "    },  # To correctly wrap around the year, we need to specify the resampling frequency.\n",
    "    \"august\": {\"month\": [8]},\n",
    "    \"annual\": {},\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### b) Defining criteria for missing values\n",
    "\n",
    "The `xhydro.indicators.get_yearly_op` function also includes two arguments, `missing` and `missing_options`, which allow you to define tolerances for missing data. These arguments leverage `xclim` to handle missing values, and the available options are detailed in the [xclim documentation](https://xclim.readthedocs.io/en/stable/checks.html#missing-values-identification).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xclim\n",
    "\n",
    "print(xclim.core.missing.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### c) Simple example\n",
    "\n",
    "Let's start with a straightforward example using the `timeargs` dictionary defined earlier. In this case, we will set a tolerance where a maximum of 15% of missing data in a year will be considered acceptable for the year to be regarded as valid.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa = xh.indicators.get_yearly_op(\n",
    "    ds, op=\"max\", timeargs=timeargs, missing=\"pct\", missing_options={\"tolerance\": 0.15}\n",
    ")\n",
    "\n",
    "ds_4fa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa.q_max_summer.dropna(\"time\", how=\"all\").hvplot(\n",
    "    x=\"time\", grid=True, widget_location=\"bottom\", groupby=\"id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### d) Advanced example: Using custom seasons per year or per station\n",
    "\n",
    "Customizing date ranges for each year or each station is not directly supported by `xhydro.indicators.get_yearly_op`. However, users can work around this limitation by masking their data before calling the function. When applying this approach, be sure to adjust the `missing` argument to accommodate the changes in the data availability.\n",
    "\n",
    "In this example, we will define a season that starts on a random date in April and ends on a random date in June. Since we will mask almost the entire year, the tolerance for missing data should be adjusted accordingly. Instead of setting a general tolerance for missing data, we will use the `at_least_n` method to specify that at least 45 days of data must be available for the period to be considered valid.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nyears = np.unique(ds.time.dt.year).size\n",
    "dom_start = xr.DataArray(\n",
    "    np.random.randint(1, 30, size=(nyears,)).astype(\"str\"),\n",
    "    dims=(\"year\"),\n",
    "    coords={\"year\": np.unique(ds.time.dt.year)},\n",
    ")\n",
    "dom_end = xr.DataArray(\n",
    "    np.random.randint(1, 30, size=(nyears,)).astype(\"str\"),\n",
    "    dims=(\"year\"),\n",
    "    coords={\"year\": np.unique(ds.time.dt.year)},\n",
    ")\n",
    "\n",
    "mask = xr.zeros_like(ds[\"q\"])\n",
    "for y in dom_start.year.values:\n",
    "    # Random mask of dates per year, between April and June.\n",
    "    mask.loc[\n",
    "        {\n",
    "            \"time\": slice(\n",
    "                str(y) + \"-04-\" + str(dom_start.sel(year=y).item()),\n",
    "                str(y) + \"-06-\" + str(dom_end.sel(year=y).item()),\n",
    "            )\n",
    "        }\n",
    "    ] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask.hvplot(x=\"time\", grid=True, widget_location=\"bottom\", groupby=\"id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The name of the indexer will be used to identify the variable created here\n",
    "timeargs_custom = {\"custom\": {}}\n",
    "\n",
    "# We use .where() to mask the data that we want to ignore\n",
    "masked = ds.where(mask == 1)\n",
    "# Since we masked almost all the year, our tolerance for missing data should be changed accordingly\n",
    "missing = \"at_least_n\"\n",
    "missing_options = {\"n\": 45}\n",
    "\n",
    "# We use xr.merge() to combine the results with the previous dataset.\n",
    "ds_4fa = xr.merge(\n",
    "    [\n",
    "        ds_4fa,\n",
    "        xh.indicators.get_yearly_op(\n",
    "            masked,\n",
    "            op=\"max\",\n",
    "            timeargs=timeargs_custom,\n",
    "            missing=missing,\n",
    "            missing_options=missing_options,\n",
    "        ),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa.q_max_custom.dropna(\"time\", how=\"all\").hvplot(\n",
    "    x=\"time\", grid=True, widget_location=\"bottom\", groupby=\"id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### e) Alternative variable: Computing volumes\n",
    "\n",
    "Frequency analysis can also be applied to volumes, following a similar workflow to that of streamflow data. The main difference is that if we're starting with streamflow data, we must first convert it into volumes using `xhydro.indicators.compute_volume` (e.g., going from `m3 s-1` to `m3`). Additionally, if necessary, the `get_yearly_op` function includes an argument, `interpolate_na`, which can be used to interpolate missing data before calculating the sum.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a daily volume from a daily streamflow\n",
    "ds[\"volume\"] = xh.indicators.compute_volume(ds[\"q\"], out_units=\"hm3\")\n",
    "\n",
    "# We'll take slightly different indexers\n",
    "timeargs_vol = {\"spring\": {\"date_bounds\": [\"04-30\", \"06-15\"]}, \"annual\": {}}\n",
    "\n",
    "# The operation that we want here is the sum, not the max.\n",
    "ds_4fa = xr.merge(\n",
    "    [\n",
    "        ds_4fa,\n",
    "        xh.indicators.get_yearly_op(\n",
    "            ds,\n",
    "            op=\"sum\",\n",
    "            input_var=\"volume\",\n",
    "            timeargs=timeargs_vol,\n",
    "            missing=\"pct\",\n",
    "            missing_options={\"tolerance\": 0.15},\n",
    "            interpolate_na=True,\n",
    "        ),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa.volume_sum_spring.dropna(\"time\", how=\"all\").hvplot(\n",
    "    x=\"time\", grid=True, widget_location=\"bottom\", groupby=\"id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local frequency analysis\n",
    "\n",
    "After extracting the raw data, such as annual maximums or minimums, the local frequency analysis is performed in three steps:\n",
    "\n",
    "1. Use `xhydro.frequency_analysis.local.fit` to determine the best set of parameters for a given number of statistical distributions.\n",
    "2. (Optional) Use `xhydro.frequency_analysis.local.criteria` to compute goodness-of-fit criteria and assess how well each statistical distribution fits the data.\n",
    "3. Use `xhydro.frequency_analysis.local.parametric_quantiles` to calculate return levels based on the fitted parameters.\n",
    "\n",
    "To speed up the Notebook, we'll only perform the analysis on a subset of variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.local.fit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `fit` function enables the fitting of multiple statistical distributions simultaneously, such as `[\"genextreme\", \"pearson3\", \"gumbel_r\", \"expon\"]`. Since different distributions have varying parameter sets (and sometimes different naming conventions), `xHydro` handles this complexity by using a `dparams` dimension, filling in NaN values where needed. When the results interact with `SciPy`, such as the `parametric_quantiles` function, only the relevant parameters for each distribution are passed. The selected distributions are stored in a newly created `scipy_dist` dimension.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = xhfa.local.fit(ds_4fa[[\"q_max_spring\", \"volume_sum_spring\"]], min_years=15)\n",
    "\n",
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Criteria like AIC (Akaike Information Criterion), BIC (Bayesian Information Criterion), and AICC (Corrected AIC) are valuable tools for comparing the fit of different statistical models. These criteria balance the goodness-of-fit of a model with its complexity, helping to avoid overfitting. AIC and AICC are particularly useful when comparing models with different numbers of parameters, while BIC tends to penalize complexity more heavily, making it more conservative. Lower values of these criteria indicate better model performance, with AICC being especially helpful in small sample sizes. By using these metrics, we can objectively determine the most appropriate model for our data.\n",
    "\n",
    "These three criteria can be accessed using `xhydro.frequency_analysis.local.criteria`. The results are added to a new `criterion` dimension. In this example, the AIC, BIC, and AICC all provide a weak indication that the Generalized Extreme Value (GEV) distribution is the best fit for the data, though the Gumbel distribution may also be a suitable choice. Conversely, the Pearson III failed to converge and the exponential distribution was rejected based on these criteria, suggesting that they do not adequately fit the data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.local.criteria)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criteria = xhfa.local.criteria(ds_4fa[[\"q_max_spring\", \"volume_sum_spring\"]], params)\n",
    "\n",
    "criteria[\"q_max_spring\"].isel(id=0).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, return periods can be obtained using `xhfa.local.parametric_quantiles`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.local.parametric_quantiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rp = xhfa.local.parametric_quantiles(params, return_period=[2, 20, 100])\n",
    "\n",
    "rp.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a future release, plotting will be managed by a dedicated function. For now, we demonstrate the process using preliminary utilities in this notebook.\n",
    "\n",
    "The function `xhfa.local._prepare_plots` generates the data points required to visualize the results of the frequency analysis. If `log=True`, it will return log-spaced x-values between `xmin` and `xmax`. Meanwhile, `xhfa.local._get_plotting_positions` calculates plotting positions for all variables in the dataset. It accepts `alpha` and `beta` arguments. For typical values, refer to the [SciPy documentation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.mstats.plotting_positions.html). By default, `(0.4, 0.4)` is used, which corresponds to the quantile unbiased method (Cunnane).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = xhfa.local._prepare_plots(params, xmin=1, xmax=1000, npoints=50, log=True)\n",
    "pp = xhfa.local._get_plotting_positions(ds_4fa[[\"q_max_spring\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the distributions\n",
    "p1 = data.q_max_spring.hvplot(\n",
    "    x=\"return_period\", by=\"scipy_dist\", grid=True, groupby=[\"id\"], logx=True\n",
    ")\n",
    "\n",
    "# Plot the observations\n",
    "p2 = pp.hvplot.scatter(\n",
    "    x=\"q_max_spring_pp\",\n",
    "    y=\"q_max_spring\",\n",
    "    grid=True,\n",
    "    groupby=[\"id\"],\n",
    "    logx=True,\n",
    ")\n",
    "\n",
    "# And now combining the plots\n",
    "p1 * p2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uncertainties\n",
    "\n",
    "Uncertainties are an important aspect of frequency analysis and should be considered when interpreting results. These uncertainties often stem from data quality, the choice of distribution, and the estimation of parameters. While visualizations can provide insights into the model fit, it’s crucial to quantify and account for uncertainties, such as confidence intervals for parameter estimates, to ensure robust conclusions.\n",
    "\n",
    "In order to manage computational intensity, we will focus on a single catchment and limit the analysis to the two distributions that appeared to best fit the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa = ds_4fa.sel(id=\"020602\")[[\"q_max_spring\"]]\n",
    "params = params.sel(id=\"020602\", scipy_dist=[\"genextreme\", \"gumbel_r\"])[\n",
    "    [\"q_max_spring\"]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a) Bootstrapping the observations\n",
    "\n",
    "One method for quantifying uncertainties is to bootstrap the observations. In this example, we will perform bootstrapping a small number of times to illustrate the process, though in practice, a higher number of iterations (e.g., 5000) is recommended to obtain more reliable estimates. Bootstrapping resamples the observed data with replacement to generate multiple datasets, which can then be used to assess the variability in the model's parameters and results.\n",
    "\n",
    "This can be accomplished by calling `xhydro.frequency_analysis.uncertainties.bootstrap_obs`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.uncertainties.bootstrap_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_4fa_boot_obs = xhfa.uncertainties.bootstrap_obs(ds_4fa, n_samples=35)\n",
    "ds_4fa_boot_obs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This process will add a new dimension, `samples`, to the dataset. When used in conjunction with `xhydro.frequency_analysis.local.fit`, a new set of parameters will be computed for each sample. As a result, bootstrapping can become computationally expensive, especially as the number of bootstrap iterations increases.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_boot_obs = xhfa.local.fit(\n",
    "    ds_4fa_boot_obs, distributions=[\"genextreme\", \"gumbel_r\"]\n",
    ")\n",
    "rp_boot_obs = xhfa.local.parametric_quantiles(\n",
    "    params_boot_obs, return_period=[2, 20, 100]\n",
    ")\n",
    "rp_boot_obs.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### b) Bootstrapping the distributions\n",
    "\n",
    "In this approach, rather than resampling the observations directly, we resample the fitted distributions to estimate the uncertainty. This method allows us to assess the variability in the fitted distributions' parameters. As with the previous example, we will perform a minimal number of bootstrap iterations to reduce computational load, but in practice, a higher number of iterations would provide more robust estimates of uncertainty.\n",
    "\n",
    "This can be accomplished by calling `xhydro.frequency_analysis.uncertainties.bootstrap_dist`. Unlike `bootstrap_obs`, this method does not support lazy evaluation and requires a specific function for the fitting step: `xhydro.frequency_analysis.uncertainties.fit_boot_dist`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.uncertainties.bootstrap_dist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(xhfa.uncertainties.fit_boot_dist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_values = xhfa.uncertainties.bootstrap_dist(ds_4fa, params, n_samples=35)\n",
    "params_boot_dist = xhfa.uncertainties.fit_boot_dist(tmp_values)\n",
    "params_boot_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rp_boot_dist = xhfa.local.parametric_quantiles(\n",
    "    params_boot_dist.load(), return_period=[2, 20, 100]\n",
    ")  # Lazy computing is not supported\n",
    "rp_boot_dist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### c) Comparison\n",
    "\n",
    "Let's show the difference between both approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "fig.set_figheight(4)\n",
    "fig.set_figwidth(15)\n",
    "\n",
    "# Subset the data\n",
    "rp_orig = rp.q_max_spring.sel(id=\"020602\", scipy_dist=\"genextreme\")\n",
    "boot_obs = rp_boot_obs.q_max_spring.sel(scipy_dist=\"genextreme\")\n",
    "boot_dist = rp_boot_dist.q_max_spring.sel(scipy_dist=\"genextreme\")\n",
    "\n",
    "# Original fit\n",
    "ax.plot(\n",
    "    rp_orig.return_period.values,\n",
    "    rp_orig,\n",
    "    \"black\",\n",
    "    label=\"Original fit\",\n",
    ")\n",
    "\n",
    "ax.plot(\n",
    "    boot_obs.return_period.values,\n",
    "    boot_obs.quantile(0.5, \"samples\"),\n",
    "    \"red\",\n",
    "    label=\"Bootstrapped observations\",\n",
    ")\n",
    "boot_obs_05 = boot_obs.quantile(0.05, \"samples\")\n",
    "boot_obs_95 = boot_obs.quantile(0.95, \"samples\")\n",
    "ax.fill_between(\n",
    "    boot_obs.return_period.values, boot_obs_05, boot_obs_95, alpha=0.2, color=\"red\"\n",
    ")\n",
    "\n",
    "ax.plot(\n",
    "    boot_dist.return_period.values,\n",
    "    boot_dist.quantile(0.5, \"samples\"),\n",
    "    \"blue\",\n",
    "    label=\"Bootstrapped distribution\",\n",
    ")\n",
    "boot_dist_05 = boot_dist.quantile(0.05, \"samples\")\n",
    "boot_dist_95 = boot_dist.quantile(0.95, \"samples\")\n",
    "ax.fill_between(\n",
    "    boot_dist.return_period.values, boot_dist_05, boot_dist_95, alpha=0.2, color=\"blue\"\n",
    ")\n",
    "\n",
    "plt.xscale(\"log\")\n",
    "plt.grid(visible=True)\n",
    "plt.xlabel(\"Return period (years)\")\n",
    "plt.ylabel(\"Streamflow (m$^3$/s)\")\n",
    "ax.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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}