{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a80c1c99",
   "metadata": {},
   "source": [
    "## Changing response functions\n",
    "*R.A. Collenteur, University of Graz, 2021*\n",
    "\n",
    "In this notebook the new `ChangeModel` is tested, based on the work by [Obergfjell et al. (2019)](https://ngwa.onlinelibrary.wiley.com/doi/10.1111/gwat.12891). The main idea is to apply different response functions for two different periods. As an example we look at the the groundwater levels measured near the river the Mur in Austria, where a dam was recently built. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e6ca87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import pastas as ps\n",
    "\n",
    "ps.set_log_level(\"ERROR\")\n",
    "ps.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bf6d14e",
   "metadata": {},
   "source": [
    "## 1. Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce526019",
   "metadata": {},
   "outputs": [],
   "source": [
    "prec = pd.read_csv(\"data_step/prec.csv\", index_col=0, parse_dates=True).squeeze()\n",
    "evap = pd.read_csv(\"data_step/evap.csv\", index_col=0, parse_dates=True).squeeze()\n",
    "head = pd.read_csv(\"data_step/head.csv\", index_col=0, parse_dates=True).squeeze()\n",
    "river = pd.read_csv(\"data_step/river.csv\", index_col=0, parse_dates=True).squeeze()\n",
    "river -= river.min()\n",
    "\n",
    "axes = ps.plots.series(\n",
    "    head=head,\n",
    "    stresses=[prec, evap, river],\n",
    "    tmin=\"2000\",\n",
    "    labels=[\"Head\\n[m]\", \"Prec\\n[mm/d]\", \"Evap\\n[mm/d]\", \"River [m3/d]\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c06edf2",
   "metadata": {},
   "source": [
    "## 2. The weighting factor\n",
    "\n",
    "The stress is convoluted two times with different response functions. Then, a weighting function is used to add the two contributions together and compute the final contribution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37db861",
   "metadata": {},
   "outputs": [],
   "source": [
    "npoints = 100\n",
    "\n",
    "tchange = 50 / npoints\n",
    "t = np.linspace(0, 1, npoints)\n",
    "color = plt.cm.viridis(np.linspace(0, 1, 10))\n",
    "\n",
    "for beta, c in zip(np.linspace(-1, 1, 10), color):\n",
    "    beta1 = beta * npoints\n",
    "    omega = 1 / (np.exp(beta1 * (t - tchange)) + 1)\n",
    "    plt.plot(omega, color=c, label=\"$beta$={:.2f}\".format(beta))\n",
    "\n",
    "plt.ylabel(\"$\\omega$ [-]\")\n",
    "plt.xlabel(\"Time [t]\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d1fda0",
   "metadata": {},
   "source": [
    "## 3. Make a model\n",
    "We now make two models:\n",
    "\n",
    "- one model where we assume the response of the heads to the river level remains the same\n",
    "- and one model where the response to the river levels changes. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01c3c160",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normal Model\n",
    "ml = ps.Model(head, name=\"linear\")\n",
    "\n",
    "sm = ps.StressModel(river, ps.Exponential(), name=\"test\")\n",
    "step = ps.StepModel(\"2012-01-01\", rfunc=ps.One(), name=\"step\")\n",
    "\n",
    "ml.add_stressmodel([sm, step])\n",
    "ml.solve(report=False, tmin=\"2004\", tmax=\"2017-12-31\", noise=True)\n",
    "ml.plots.results(figsize=(10, 6))\n",
    "\n",
    "# ChangeModel\n",
    "ml2 = ps.Model(head, name=\"linear\")\n",
    "\n",
    "cm = ps.ChangeModel(\n",
    "    river,\n",
    "    ps.Exponential(),\n",
    "    ps.Exponential(),\n",
    "    name=\"test\",\n",
    "    tchange=\"2012-01-01\",\n",
    ")\n",
    "\n",
    "ml2.add_stressmodel([cm, step])\n",
    "ml2.solve(report=False, tmin=\"2004\", tmax=\"2017-12-31\", noise=True)\n",
    "ml2.plots.results(figsize=(10, 6));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "730f8336",
   "metadata": {},
   "source": [
    "The second model shows a better fit, but also the step trend changed.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaba0484",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"RMSE for the first model:\", ml.stats.rmse().round(2))\n",
    "print(\"RMSE for the second model:\", ml2.stats.rmse().round(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a0e9a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml2.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea69949",
   "metadata": {},
   "source": [
    "## 4. Compare the response functions\n",
    "We can also look at the response to the river before and after, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f7dd85",
   "metadata": {},
   "outputs": [],
   "source": [
    "cm_rf1 = cm.rfunc1.step(\n",
    "    p=ml2.parameters.loc[[\"test_1_A\", \"test_1_a\"], \"optimal\"].values\n",
    ")\n",
    "cm_rf2 = cm.rfunc2.step(\n",
    "    p=ml2.parameters.loc[[\"test_2_A\", \"test_2_a\"], \"optimal\"].values\n",
    ")\n",
    "\n",
    "plt.plot(np.arange(1, len(cm_rf1) + 1), cm_rf1)\n",
    "plt.plot(np.arange(1, len(cm_rf2) + 1), cm_rf2)\n",
    "\n",
    "plt.legend([\"Before\", \"After\"]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a04ae63",
   "metadata": {},
   "source": [
    "## 5. Another way\n",
    "We can also add the stress twice, saving one parameter that needs to be estimated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c99291",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml3 = ps.Model(head, name=\"linear\")\n",
    "\n",
    "river1 = river.copy()\n",
    "river1.loc[\"2012\":] = 0\n",
    "\n",
    "river2 = river.copy()\n",
    "river2.loc[:\"2011\"] = 0\n",
    "\n",
    "r1 = ps.StressModel(river1, rfunc=ps.Exponential(), name=\"river\")\n",
    "r2 = ps.StressModel(river2, rfunc=ps.Exponential(), name=\"river2\")\n",
    "step = ps.StepModel(\"2012-01-01\", rfunc=ps.One(), name=\"step\")\n",
    "\n",
    "ml3.add_stressmodel([r1, r2, step])\n",
    "ml3.solve(report=False, tmin=\"2004\", tmax=\"2017-12-31\", noise=True)\n",
    "ml3.plots.results(figsize=(10, 6));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa396913",
   "metadata": {},
   "source": [
    "## How do the results compare?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f9d5bfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change model\n",
    "cm_rf1 = cm.rfunc1.step(\n",
    "    p=ml2.parameters.loc[[\"test_1_A\", \"test_1_a\"], \"optimal\"].values\n",
    ")\n",
    "cm_rf2 = cm.rfunc2.step(\n",
    "    p=ml2.parameters.loc[[\"test_2_A\", \"test_2_a\"], \"optimal\"].values\n",
    ")\n",
    "plt.plot(np.arange(1, len(cm_rf1) + 1), cm_rf1)\n",
    "plt.plot(np.arange(1, len(cm_rf2) + 1), cm_rf2)\n",
    "\n",
    "# 2 stressmodels\n",
    "ml3.get_step_response(\"river\").plot()\n",
    "ml3.get_step_response(\"river2\").plot()\n",
    "\n",
    "plt.legend(\n",
    "    [\n",
    "        \"Before (ChangeModel)\",\n",
    "        \"After (ChangeModel)\",\n",
    "        \"Before (method 2)\",\n",
    "        \"After (method 2)\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a071f27d",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "Obergfell, C., Bakker, M. and Maas, K. (2019), Identification and Explanation of a Change in the Groundwater Regime using Time Series Analysis. Groundwater, 57: 886-894. https://doi.org/10.1111/gwat.12891"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pastas_dev",
   "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "29475f8be425919747d373d827cb41e481e140756dd3c75aa328bf3399a0138e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}