{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Threshold non-linearities\n",
    "*Developed by Ruben Caljé*\n",
    "\n",
    "This notebook compares two different options in Pastas for modeling threshold non-linear groundwater systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import pastas as ps\n",
    "\n",
    "ps.set_log_level(\"WARNING\")\n",
    "ps.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start with a basic model that contains a `RechargeModel` to model the influence of precipitation and evaporation on groundwater head. We can see that the simulation of this model does not reproduce the valleys and peaks of the measurement-series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the data\n",
    "kwargs = dict(index_col=0, parse_dates=[0])\n",
    "head = pd.read_csv(\"data_notebook_8/B28H1804_2.csv\", **kwargs).squeeze(\"columns\")\n",
    "prec = pd.read_csv(\"data_notebook_8/prec.csv\", **kwargs).squeeze(\"columns\")\n",
    "evap = pd.read_csv(\"data_notebook_8/evap.csv\", **kwargs).squeeze(\"columns\")\n",
    "\n",
    "# generate a model and solve\n",
    "ml = ps.Model(head)\n",
    "ml.add_stressmodel(ps.RechargeModel(prec, evap))\n",
    "ml.solve()\n",
    "\n",
    "# and we plot the results\n",
    "ml.plots.results();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ThresholdTransform\n",
    "We can add a ThresholdTransform to model a threshold above which the groundwater reaction is damped. This transform is applied after the simulation is calculated. Therefore it can be added to any model. It adds two extra parameters: the level above and the factor by which the groundwater levels are damped. It is very effective for simulating the selected groundwater series. The $R^2$ value increases from 87% to 96%. We can see that the fit with the observations is much better, for both low and high measurement values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.add_transform(ps.ThresholdTransform())\n",
    "ml.solve(report=False)\n",
    "ml.plots.results();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TarsoModel\n",
    "We can also model this series using a TarsoModel. Tarso stands for Threshold AutoRegressive Self-exciting Open-loop. The simulation is calculated by two exponential response functions, where the second response function becomes active when the simulation reaches a certain threshold-value.\n",
    "\n",
    "Compared to the ThresholdTransform the simulation is not only damped above the threshold, but also the response time is changed above the threshold. A large drawback of the TarsoModel however is that it only allows the Exponential response function and it cannot be combined with other model elements (stressmodels, constant or transform). Therefore, all those other elements are removed from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm = ml.stressmodels[\"recharge\"]\n",
    "prec = sm.prec.series\n",
    "evap = sm.evap.series\n",
    "\n",
    "# delete all the stressmodels, the constant and the transform\n",
    "ml.del_stressmodel(\"recharge\")\n",
    "ml.del_constant()\n",
    "ml.del_transform()\n",
    "\n",
    "# then add a TarsoModel\n",
    "sm = ps.TarsoModel(prec, evap, oseries=ml.oseries.series)\n",
    "ml.add_stressmodel(sm)\n",
    "\n",
    "# and solve and plot the results again\n",
    "ml.solve(report=False)\n",
    "ml.plots.results();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit of the Tarso model (two exponential response functions) is similar to the fit of the Gamma response function with a ThresholdTransform (a damping transformation above a threshold).\n",
    "It is possible to interpret the TarsoModel as a physical model. In this model, there are two discharges with different resistances, where the second discharge is not always active. This model can be visualized by the image below (taken from https://edepot.wur.nl/406715).\n",
    "\n",
    "![Tarso system](data_notebook_8/tarso_system.png)\n",
    "\n",
    "In this model d1 and c1 are equal to our parameters d1 and A1, d2 and c2 are equal to our parameters d0 and A0. We can then calculate the water balance and plot it with the code below. In this plot, all the in-fluxes (mainly precipitation) are positive, and all the out-fluxes (mainly evaporation) are negative. An exception is the storage term, to make sure the positive an negative balance terms level out. An increase in storage has a negative sign, and a decrease in storage has a positive sign."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the water balance\n",
    "sim = ml.simulate()\n",
    "P = prec[sim.index[0] :]\n",
    "E = -evap[sim.index[0] :]\n",
    "p = ml.get_parameters(\"tarso\")\n",
    "Q0 = -(sim - p[2]) / p[0]\n",
    "Q1 = -(sim - p[5]) / p[3]\n",
    "Q1[sim < p[5]] = 0.0\n",
    "# calculate storage\n",
    "S = -(P + E + Q0 + Q1)\n",
    "# combine these Series in a DataFrame\n",
    "df = pd.DataFrame({\"P\": P, \"E\": E, \"Q0\": Q0, \"Q1\": Q1, \"S\": S}) * 1000\n",
    "# resample the balance to monthly values, to make the graph more readable\n",
    "df = df.resample(\"ME\").mean()\n",
    "# and set the index to the middle of the month\n",
    "df.index = df.index - (df.index - (df.index - pd.offsets.MonthBegin())) / 2\n",
    "\n",
    "# make a new figure\n",
    "f, ax = plt.subplots(nrows=2, sharex=True, figsize=(14, 8), layout=\"constrained\")\n",
    "\n",
    "# plot heads in the upper graph\n",
    "ax[0].set_ylabel(\"Groundwater head (m to MSL)\")\n",
    "sim.plot(ax=ax[0], x_compat=True)\n",
    "ml.observations().plot(\n",
    "    ax=ax[0], marker=\".\", color=\"k\", x_compat=True, markersize=2, linestyle=\"none\"\n",
    ")\n",
    "ax[0].axhline(p[2], linestyle=\"--\", color=\"C2\")\n",
    "ax[0].axhline(p[5], linestyle=\"--\", color=\"C3\")\n",
    "\n",
    "# plot discharges in the lower graph\n",
    "ax[1].set_ylabel(\"Monthly averaged flow rate (mm/d)\")\n",
    "color = [\"C0\", \"C1\", \"C2\", \"C3\", \"C4\"]\n",
    "\n",
    "df_up = df.where(df > 0, np.nan)\n",
    "df_down = df.where(df < 0, np.nan)\n",
    "df_up.plot.area(ax=ax[1], color=color, linewidth=0)\n",
    "ax[1].set_ylim(-1, 1)  # set ylim so we see no warning\n",
    "df_down.plot.area(ax=ax[1], color=color, linewidth=0, legend=False)\n",
    "ax[1].axhline(0.0, linestyle=\"--\", color=\"k\")\n",
    "\n",
    "# set some stuff for both axes\n",
    "for iax in ax:\n",
    "    iax.autoscale(tight=True)\n",
    "    iax.minorticks_off()\n",
    "    iax.grid(True)"
   ]
  }
 ],
 "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.12.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "29475f8be425919747d373d827cb41e481e140756dd3c75aa328bf3399a0138e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}