{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65ea179b",
   "metadata": {},
   "source": [
    "# Checking Pastas models\n",
    "\n",
    "*Developed by D.A. Brakenhoff, Artesia*\n",
    "\n",
    "This notebooks showcases the `pastas.check` submodule. This module can be used to check your time series models with pre-defined checks, or users can define their own checks.\n",
    "\n",
    "For more elaborate discussions on why checking different aspects of your model is a good idea, we refer you to some of the other notebooks in the documentation, e.g. [diagnostic_checking.ipynb](https://pastas.readthedocs.io/latest/examples/diagnostic_checking.html) and/or [stowa_calibration.ipynb](https://pastas.readthedocs.io/latest/examples/stowa_calibration.html). In this notebook we focus on the `pastas.check` tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "651d407a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "import pastas as ps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "684b4f6a",
   "metadata": {},
   "source": [
    "## Build a simple model\n",
    "\n",
    "First let's load some head data and build a simple time series model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09598842",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = pd.read_csv(\"data/head_nb1.csv\", index_col=0, parse_dates=True).squeeze(\"columns\")\n",
    "rain = pd.read_csv(\"data/rain_nb1.csv\", index_col=0, parse_dates=True).squeeze(\n",
    "    \"columns\"\n",
    ")\n",
    "evap = pd.read_csv(\"data/evap_nb1.csv\", index_col=0, parse_dates=True).squeeze(\n",
    "    \"columns\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4430cc63",
   "metadata": {},
   "source": [
    "We use a linear recharge model, so $R = P - f \\cdot E$. Build the model and solve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae3b22e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = ps.Model(obs, name=\"groundwater_head\")\n",
    "sm = ps.RechargeModel(prec=rain, evap=evap, rfunc=ps.Gamma(), name=\"recharge\")\n",
    "ml.add_stressmodel(sm)\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f520fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = ml.plot(figsize=(10, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ce730b3",
   "metadata": {},
   "source": [
    "## Checking the model\n",
    "\n",
    "The pastas `check` module contains check functions and a convenience function if you\n",
    "want to do multiple checks on a single model.\n",
    "\n",
    "Let's inspect the module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dec4ad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ed55260",
   "metadata": {},
   "source": [
    "### What is a check function?\n",
    "\n",
    "Check functions are the methods that determine whether a pastas Model meets certain\n",
    "criteria (passes the check) or fails to meet those criteria (fails the check). The\n",
    "criteria can be anything, from goodness-of-fit statistics to the magnitude of the\n",
    "variation of some contribution to the head. It is up to the modeller to decide which\n",
    "criteria their model should be subjected to.\n",
    "\n",
    "Check functions follow a specific template:\n",
    "\n",
    " - A `pastas.Model` as the first argument\n",
    " - Any number of keyword arguments after that\n",
    " - Return a DataFrame containing information about the performed check. The following columns are defined:\n",
    "    - **statistic:** the test statistic\n",
    "    - **operator:** the operator, e.g. >, ==, within, etc.\n",
    "    - **threshold:** the user-specified threshold or comparison values\n",
    "    - **dimensions:** units or dimensions of statistic (useful for interpretation)\n",
    "    - **pass:** whether the check is passed or not\n",
    "    - **comment:** any additional comments relevant to the check"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae5145ab",
   "metadata": {},
   "source": [
    "### Pre-defined checks\n",
    "\n",
    "The `pastas.check` module contains several pre-defined checks that are frequently\n",
    "applied to time series models. These can be listed with `ps.check.checks` which is a\n",
    "dictionary containing the names and functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69d20afd",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(ps.check.checks.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89c9be46",
   "metadata": {},
   "source": [
    "#### Check: $R^2 \\geq$ threshold\n",
    "\n",
    "Let's try applying the `rsq_geq_threshold` check to our model, which tests whether the\n",
    "$R^2 \\geq s$, where $s$ is some user-defined threshold. We can already see the $R^2$\n",
    "fit in the plot of the model simulation above, so let's say we want our fit to be\n",
    "greater or equal to 0.9. The model should obviously pass this check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28b617e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check.rsq_geq_threshold(ml, threshold=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aacd509c",
   "metadata": {},
   "source": [
    "This example isn't very interesting, but you can already\n",
    "see that when we want to apply multiple checks to all of our models, having each of\n",
    "these checks contained within a function can be useful."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b54ea69",
   "metadata": {},
   "source": [
    "### Writing your own checks\n",
    "\n",
    "Now as an extra example let's write our own check function. This time we want to check\n",
    "whether our model is sufficiently good at simulating the low groundwater levels in\n",
    "summer.\n",
    "\n",
    "We'll do this by checking the goodness-of-fit for the summer periods. We can write a\n",
    "function that accepts a list of months in which to consider the residuals in order to\n",
    "calculate the fit statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e71839d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rsq_geq_threshold_in_months(\n",
    "    ml: ps.Model, threshold: float, months: list[int]\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"Check if the R² of the model is >= to a threshold in specific months.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    ml : pastas.Model\n",
    "        The Pastas model to check.\n",
    "    threshold : float\n",
    "        The R² threshold value.\n",
    "    months : list of ints\n",
    "        The month numbers to consider in each year.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    pd.DataFrame\n",
    "        A DataFrame showing the check results\n",
    "    \"\"\"\n",
    "    res = ml.residuals()  # get the model residuals\n",
    "    mask = res.index.month.isin(months)  # define a mask for the selected months\n",
    "    rsq = ps.stats.rsq(obs=ml.observations().loc[mask], res=res.loc[mask])  # compute R²\n",
    "    # store results and context\n",
    "    df = ps.check.get_empty_check_dataframe()\n",
    "    df.loc[\"rsq_geq_threshold_in_months\"] = (\n",
    "        rsq,\n",
    "        \">=\",\n",
    "        threshold,\n",
    "        \"-\",\n",
    "        rsq >= threshold,\n",
    "        f\"in months: {months}\",\n",
    "    )\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7fb329",
   "metadata": {},
   "outputs": [],
   "source": [
    "rsq_geq_threshold_in_months(ml, threshold=0.8, months=[6, 7, 8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05aede01",
   "metadata": {},
   "source": [
    "## Applying multiple checks to a model\n",
    "\n",
    "Often we want to check multiple criteria when deciding whether our model is fit for\n",
    "purpose. \n",
    "\n",
    "### Pre-defined checks\n",
    "\n",
    "Some of the pre-defined checks already check multiple criteria, for example whether\n",
    "the parameters do not lie on the parameter bounds. These methods return a DataFrame\n",
    "with multiple rows, with the results of each check on a separate row.\n",
    "\n",
    "#### Check: Parameter bounds\n",
    "\n",
    "This check is automatically performed when solving a time series model, and any warnings are logged and reported at the bottom of the fit report. But it can be useful to include it in your list of checks. When an optimal parameter lies on a boundary, the check fails."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24200bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check.parameter_bounds(ml)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9caa1bd9",
   "metadata": {},
   "source": [
    "#### Check: Uncertainty parameters\n",
    "\n",
    "This check tests whether the absolute value of the optimal parameters is larger than some factor times the estimated standard deviation of the parameter. Basically, it checks whether the parameter can be estimated with sufficient certainty. This check requires that the estimate of $\\sigma$ is reliable, which is the case when the noise meets certain requirements. This check always includes this comment, to remind users to check the noise prior to trusting the estimated parameter uncertainties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6efe0a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check.uncertainty_parameters(ml, n_std=1.96)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54a75724",
   "metadata": {},
   "source": [
    "#### Check: Length response relative to calibratiod period or warmup\n",
    "\n",
    "The memory of the response function indicates the time it takes until the effect of some change has taken place. For example, the time it takes for the groundwater level to rise after 1 mm of rain today. This time is commonly expressed as a percentage e.g. $t_{95}$, the time it takes until 95% of the rise has taken place. (Since the response is asymptotic, the 100% lies at time infinitiy.)\n",
    "\n",
    "This memory should never be longer than the calibration period, and preferably be significantly shorter. By default the $t_{95}$ is used and it is compared to half the length of the calibration period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2d87926",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check.response_memory(ml, cutoff=0.95, factor_length_oseries=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "436ffcf9",
   "metadata": {},
   "source": [
    "The memory can also compared to the warmup period. When the memory exceeds the warmup period, the model is not yet done warming up by the time the simulation starts, which can cause issues when calibrating or simulating the model for different time periods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0c35386",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps.check.response_memory_vs_warmup(ml, cutoff=0.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf937090",
   "metadata": {},
   "source": [
    "### The `checklist` function for performing multiple checks\n",
    "\n",
    "The function `ps.check.checklist` is a convenience function for applying multiple\n",
    "checks to a model. This function accepts a list of checks. This list can consist of the\n",
    "following items:\n",
    "   \n",
    "   - the name of a built-in check, e.g. \"rsq_geq_threshold\"\n",
    "   - any function that requires only a model as its input\n",
    "   - a dictionary containing any function and any additional arguments to be passed to the check function.\n",
    "\n",
    "Note that in the first two cases it is not possible to alter any additional arguments\n",
    "to the functions. If relevant it will use default values. Additionally, each check\n",
    "function should return a DataFrame with the expected columns, so the results can be\n",
    "combined.\n",
    "\n",
    "The example below show-cases the application of 3 checks in each of the ways described\n",
    "above. Optionally, a report is shown (basically the resulting DataFrame) in which the\n",
    "checks (pass/fail) are colored for quick visual inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f764dd6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add checks to a list\n",
    "checklist = [\n",
    "    # name of a built-in check\n",
    "    \"rsq_geq_threshold\",\n",
    "    # any check function that only requires the model as input\n",
    "    ps.check.parameter_bounds,\n",
    "    # check function, with additional arguments\n",
    "    {\n",
    "        \"func\": rsq_geq_threshold_in_months,\n",
    "        \"threshold\": 0.8,\n",
    "        \"months\": [6, 7, 8],\n",
    "    },\n",
    "]\n",
    "\n",
    "# perform all checks on a model, and optionally display a report\n",
    "checks = ps.check.checklist(ml, checklist, report=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b96249a9",
   "metadata": {},
   "source": [
    "## Checks in literature\n",
    "\n",
    "Currently, the `ps.check` module contains two pre-defined checklists taken from two articles on time series analysis: \n",
    "\n",
    "- the checks used in Brakenhoff et al. (2022).\n",
    "- the checks used in Zaadnoordijk et al. (2019).\n",
    "\n",
    "These articles used certain checks to determine whether a time series model was\n",
    "reliable or not. Note that these checklists were developed to address specific research\n",
    "questions and hydrogeological contexts. They are provided as reference implementations\n",
    "but may not be universally applicable across all modeling scenarios.\n",
    "\n",
    "These checklists can be passed to the `ps.check.checklist()` function. Note that some\n",
    "of these checks require the noise to meet the requirement of white noise. Our model\n",
    "currently does not have a noise model, so let's see what that does to the checks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97121373",
   "metadata": {},
   "outputs": [],
   "source": [
    "checklist_literature = ps.check.get_checks_literature(\"brakenhoff_2022\")\n",
    "checklist_literature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b256a33e",
   "metadata": {},
   "outputs": [],
   "source": [
    "checks_no_noise = ps.check.checklist(ml, checklist_literature)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ead9819",
   "metadata": {},
   "source": [
    "Now let's add a noisemodel, solve the model again and run the checks again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b778a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.add_noisemodel(ps.ArNoiseModel())\n",
    "ml.solve(initial=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c92ab39",
   "metadata": {},
   "source": [
    "Note that the autocorrelation check now passes, which means we are more confident that the estimated parameter uncertainties are reliable.\n",
    "\n",
    "*__Note:__ No significant autocorrelation is just one of the criteria the noise has to meet in order to meet the requirements of white noise, but we will not discuss that here. For more information refer to the notebooks [diagnostic_checking.ipynb](https://pastas.readthedocs.io/latest/examples/diagnostic_checking.html) and/or [stowa_calibration.ipynb](https://pastas.readthedocs.io/latest/examples/stowa_calibration.html).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecf3131b",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_w_noise = ps.check.checklist(ml, checklist_literature)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e010eed9",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "- Brakenhoff DA, Vonk MA, Collenteur RA, Van Baar M and Bakker M (2022) Application of Time Series Analysis to Estimate Drawdown From Multiple Well Fields. Front. Earth Sci. 10:907609. [doi: 10.3389/feart.2022.907609](https://www.frontiersin.org/journals/earth-science/articles/10.3389/feart.2022.907609/full)\n",
    "- Zaadnoordijk, W.J., Bus, S.A.R., Lourens, A. and Berendrecht, W.L. (2019), Automated Time Series Modeling for Piezometers in the National Database of the Netherlands. Groundwater, 57: 834-843. https://doi.org/10.1111/gwat.12819"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pastas (3.13.11)",
   "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}