{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding multiple wells\n",
    "\n",
    "This notebook shows how a WellModel can be used to fit multiple wells with one response function. The influence of the individual wells is scaled by the distance to the observation point. \n",
    "\n",
    "*Developed by R.C. Caljé, (Artesia Water 2020), D.A. Brakenhoff, (Artesia Water 2019), and R.A. Collenteur, (Artesia Water 2018)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pastas as ps\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ps.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and set data\n",
    "\n",
    "Set the coordinates of the extraction wells and calculate the distances to the observation well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify coordinates observations\n",
    "xo = 85850\n",
    "yo = 383362\n",
    "\n",
    "# Specify coordinates extractions\n",
    "relevant_extractions = {\n",
    "    \"Extraction_2\": (83588, 383664),\n",
    "    \"Extraction_3\": (88439, 382339),\n",
    "}\n",
    "\n",
    "# calculate distances\n",
    "distances = []\n",
    "for extr, xy in relevant_extractions.items():\n",
    "    xw = xy[0]\n",
    "    yw = xy[1]\n",
    "    distances.append(np.sqrt((xo - xw) ** 2 + (yo - yw) ** 2))\n",
    "\n",
    "df = pd.DataFrame(\n",
    "    distances,\n",
    "    index=relevant_extractions.keys(),\n",
    "    columns=[\"Distance to observation well\"],\n",
    ")\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read the stresses from their csv files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read oseries\n",
    "oseries = pd.read_csv(\n",
    "    \"data_notebook_10/Observation_well.csv\", index_col=0, parse_dates=[0]\n",
    ").squeeze()\n",
    "oseries.name = oseries.name.replace(\" \", \"_\")\n",
    "# read stresses\n",
    "stresses = {}\n",
    "for fname in os.listdir(\"data_notebook_10\"):\n",
    "    series = pd.read_csv(\n",
    "        os.path.join(\"data_notebook_10\", fname), index_col=0, parse_dates=[0]\n",
    "    ).squeeze()\n",
    "    stresses[fname.strip(\".csv\").replace(\" \", \"_\")] = series"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then plot the observations, together with the diferent stresses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot timeseries\n",
    "f1, axarr = plt.subplots(len(stresses.keys()) + 1, sharex=True, figsize=(10, 8))\n",
    "oseries.plot(ax=axarr[0], color=\"k\")\n",
    "axarr[0].set_title(oseries.name)\n",
    "\n",
    "for i, name in enumerate(stresses.keys(), start=1):\n",
    "    stresses[name].plot(ax=axarr[i])\n",
    "    axarr[i].set_title(name)\n",
    "plt.tight_layout(pad=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a model with a separate StressModel for each extraction\n",
    "\n",
    "First we create a model with a separate StressModel for each groundwater extraction. First we create a model with the heads timeseries and add recharge as a stress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model\n",
    "ml = ps.Model(oseries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the precipitation and evaporation timeseries and round the index to remove the hours from the timestamps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prec = stresses[\"Precipitation\"]\n",
    "prec.index = prec.index.round(\"D\")\n",
    "prec.name = \"prec\"\n",
    "evap = stresses[\"Evaporation\"]\n",
    "evap.index = evap.index.round(\"D\")\n",
    "evap.name = \"evap\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a recharge stressmodel and add to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rm = ps.RechargeModel(prec, evap, ps.Exponential(), \"Recharge\")\n",
    "ml.add_stressmodel(rm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Modify the extraction timeseries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extraction_ts = {}\n",
    "\n",
    "for name in relevant_extractions.keys():\n",
    "    # get extraction timeseries\n",
    "    s = stresses[name]\n",
    "\n",
    "    # convert index to end-of-month timeseries\n",
    "    s.index = s.index.to_period(\"M\").to_timestamp(\"M\")\n",
    "\n",
    "    # resample to daily values\n",
    "    new_index = pd.date_range(s.index[0], s.index[-1], freq=\"D\")\n",
    "    s_daily = ps.ts.timestep_weighted_resample(s, new_index, fast=True).dropna()\n",
    "    name = name.replace(\" \", \"_\")\n",
    "    s_daily.name = name\n",
    "\n",
    "    # append to stresses list\n",
    "    extraction_ts[name] = s_daily"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add each of the extractions as a separate StressModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, stress in extraction_ts.items():\n",
    "    sm = ps.StressModel(stress, ps.Hantush(), name, up=False, settings=\"well\")\n",
    "    ml.add_stressmodel(sm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the results\n",
    "Plot the decomposition to see the individual influence of each of the wells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ml.plots.decomposition();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can calculate the gain of each extraction (quantified as the effect on the groundwater level of a continuous extraction of ~1 Mm$^3$/yr)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in relevant_extractions.keys():\n",
    "    sm = ml.stressmodels[name]\n",
    "    p = ml.get_parameters(name)\n",
    "    gain = sm.rfunc.gain(p) * 1e6 / 365.25\n",
    "    print(f\"{name}: gain = {gain:.3f} m / Mm^3/year\")\n",
    "    df.at[name, \"gain StressModel\"] = gain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a model with a WellModel\n",
    "We can reduce the number of parameters in the model by including the three extractions in a WellModel. This WellModel takes into account the distances from the three extractions to the observation well, and assumes constant geohydrological properties. All of the extractions now share the same response function, scaled by the distance between the extraction well and the observation well.\n",
    "\n",
    "First we create a new model and add recharge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_wm = ps.Model(oseries, oseries.name + \"_wm\")\n",
    "rm = ps.RechargeModel(prec, evap, ps.Gamma(), \"Recharge\")\n",
    "ml_wm.add_stressmodel(rm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have all the information we need to create a WellModel:\n",
    "- timeseries for each of the extractions, these are passed as a list of stresses\n",
    "- distances from each extraction to the observation point, note that the order of these distances must correspond to the order of the stresses.\n",
    "\n",
    "Note: the WellModel only works with a special version of the Hantush response function called `HantushWellModel`. This is because the response function must support scaling by a distance $r$. The HantushWellModel response function has been modified to support this. The Hantush response normally takes three parameters: the gain $A$, $a$ and $b$. This special version accepts 4 parameters: it interprets that fourth parameter as the distance $r$, and uses it to scale the parameters accordingly. \n",
    "\n",
    "Create the WellModel and add to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = ps.WellModel(list(extraction_ts.values()), \"Wells\", distances)\n",
    "ml_wm.add_stressmodel(w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the model. \n",
    "\n",
    "As we can see, the fit with the measurements (EVP) is similar to the result with the previous model, with each well included separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_wm.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the results\n",
    "Plot the decomposition to see the individual influence of each of the wells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_wm.plots.decomposition();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the stacked influence of each of the individual extraction wells in the results plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_wm.plots.stacked_results(figsize=(10, 8));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get parameters for each well (including the distance) and calculate the gain. The WellModel reorders the stresses from closest to the observation well, to furthest from the observation well. We have take this into account during the post-processing.\n",
    "\n",
    "The gain of extraction 3 is lower than the gain of extraction 2. This will always be the case in a WellModel when the distance from the observation well to extraction 3 is larger than the distance to extraction 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wm = ml_wm.stressmodels[\"Wells\"]\n",
    "for i, name in enumerate(relevant_extractions.keys()):\n",
    "    # get parameters\n",
    "    p = wm.get_parameters(model=ml_wm, istress=i)\n",
    "    # calculate gain\n",
    "    gain = wm.rfunc.gain(p) * 1e6 / 365.25\n",
    "    name = wm.stress[i].name\n",
    "    print(f\"{name}: gain = {gain:.3f} m / Mm^3/year\")\n",
    "    df.at[name, \"gain WellModel\"] = gain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate gain as function of radial distance for and plot the result, including the estimated uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.logspace(3, 3.6, 101)\n",
    "\n",
    "# calculate gain and std error vs distance\n",
    "params = ml_wm.get_parameters(wm.name)\n",
    "gain_wells = wm.rfunc.gain(params, r=wm.distances.values) * 1e6 / 365.25\n",
    "gain_vs_dist = wm.rfunc.gain(params, r=r) * 1e6 / 365.25\n",
    "gain_std_vs_dist = np.sqrt(wm.variance_gain(ml_wm, r=r)) * 1e6 / 365.25\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
    "ax.plot(r, gain_vs_dist, color=\"C0\", label=\"gain\")\n",
    "ax.plot(\n",
    "    wm.distances,\n",
    "    gain_wells,\n",
    "    color=\"C3\",\n",
    "    marker=\"o\",\n",
    "    mfc=\"none\",\n",
    "    label=\"wells\",\n",
    "    ls=\"none\",\n",
    ")\n",
    "ax.fill_between(\n",
    "    r,\n",
    "    gain_vs_dist - 2 * gain_std_vs_dist,\n",
    "    gain_vs_dist + 2 * gain_std_vs_dist,\n",
    "    alpha=0.35,\n",
    "    label=\"CI 95%\",\n",
    ")\n",
    "ax.axhline(0.0, linestyle=\"dashed\", color=\"k\")\n",
    "ax.legend(loc=(0, 1), frameon=False, ncol=3)\n",
    "ax.grid(visible=True)\n",
    "ax.set_xlabel(\"radial distance [m]\")\n",
    "ax.set_ylabel(\"gain [m / (Mm$^3$/yr)]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare individual StressModels and WellModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the gains that were calculated by the individual StressModels and the WellModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.style.format(\"{:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visually compare the two models, including the calculated contribution of the wells:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# give models descriptive name\n",
    "ml.name = \"2_wells\"\n",
    "ml_wm.name = \"WellModel\"\n",
    "\n",
    "# plot well stresses together on same plot:\n",
    "smdict = {0: [\"Recharge\"], 1: [\"Extraction_2\", \"Extraction_3\", \"Wells\"]}\n",
    "\n",
    "# comparison plot\n",
    "mc = ps.CompareModels([ml, ml_wm])\n",
    "mosaic = mc.get_default_mosaic(n_stressmodels=2)\n",
    "mc.initialize_adjust_height_figure(mosaic=mosaic, smdict=smdict)\n",
    "mc.plot(smdict=smdict)\n",
    "\n",
    "sumwells = ml.get_contribution(\"Extraction_2\") + ml.get_contribution(\"Extraction_3\")\n",
    "mc.axes[\"con1\"].plot(\n",
    "    sumwells.index, sumwells, ls=\"dashed\", color=\"C0\", label=\"sum 2_wells\"\n",
    ")\n",
    "mc.axes[\"con1\"].legend(loc=(0, 1), frameon=False, ncol=4);"
   ]
  }
 ],
 "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.9.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "dace5e1b41a98a8e52d2a8eebc3b981caf2c12e7a76736ebfb89a489e3b62e79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}