{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding wells with the Hantush response function\n",
    "\n",
    "This notebook compares the two implementations of the Hantush response function in Pastas.\n",
    "\n",
    "*Developed by D.A. Brakenhoff (Artesia, 2021)*\n",
    "\n",
    "\n",
    "## Contents\n",
    "\n",
    "- [Hantush versus HantushWellModel](#Hantush-versus-HantushWellModel)\n",
    "- [Which Hantush should I use?](#Which-Hantush-should-I-use?)\n",
    "- [Options](#Options)\n",
    "- [Synthetic example](#Synthetic-example)"
   ]
  },
  {
   "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.show_versions()\n",
    "ps.logger.setLevel(\"WARNING\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hantush versus HantushWellModel\n",
    "\n",
    "There are two implementations of the Hantush response functions in Pastas. The two implementations are very similar, but they differ in their intended application and their definition of the parameters. The table below shows the formulas for both implementations.\n",
    "\n",
    "\n",
    "| Name             | Fitting parameters  | Formula                                                                 | Description                                                                    |\n",
    "|------------------|-------------|:------------------------------------------------------------------------|--------------------------------------------------------------------------------|\n",
    "| Hantush          | 3 - A, a, b | $$ \\theta(t) = \\frac{A}{2t \\text{K}_0 \\left(2 \\sqrt{b} \\right)} e^{-t/a - ab/t} $$                               | Response function commonly used for groundwater abstraction wells.             |\n",
    "| HantushWellModel | 3 - A', a, b' | $$ \\theta(r,t) = \\frac{A^\\prime}{2t} e^{-t/a - a r^2 \\exp (b^\\prime) /t} $$ | Implementation of the Hantush well function that allows scaling with distance. |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hantush\n",
    "The Hantush response function is intended for the simulation of the effect of a\n",
    "single pumping well. The Hantush implementation has three parameters: $A$, $a$,\n",
    "and $b$. The parameter $A$ is also known as the \"gain\", which is equal to the\n",
    "steady-state contribution of a stress with unit 1. For example, the drawdown\n",
    "caused by a well with a continuous extraction rate of 1.0 (the units are\n",
    "determined by the units of the stress and head used in the model).\n",
    "\n",
    "The relationship between the parameters $A$, $a$, and $b$ and the physical\n",
    "parameters of the classic Hantush function are given in the notebook on\n",
    "[response functions](https://pastas.readthedocs.io/en/latest/concepts/response_functions.ipynb.html#Hantush-step-function-compared-to-classic-Hantush-function)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HantushWellModel\n",
    "\n",
    "The HantushWellModel also has three parameters: $A^\\prime$, $a$, and\n",
    "$b^\\prime$. The HantushWellModel response function includes the distance $r$\n",
    "between an extraction well and an observation well, which must be defined by\n",
    "the user as an input variable. This allows multiple wells to have the same\n",
    "response function, scaled by the distance $r$, which can be useful to reduce\n",
    "the number of parameters in a model with multiple extraction wells. Note that\n",
    "$r$ is a variable that must be provided by the user and is not a parameter that\n",
    "is optimized. The gain of the HantushWellModel function is\n",
    "\n",
    "$$\n",
    "A^\\prime \\text{K}_0 \\left( 2 r \\exp \\left( \\frac{b^\\prime}{2} \\right) \\right)\n",
    "$$\n",
    "\n",
    "The relationship between the parameters of the Hantush function and the\n",
    "HantushWellModel function are:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "A &= A^\\prime \\text{K}_0 \\left( 2 r \\exp \\left( \\frac{b^\\prime}{2} \\right) \\right)\\\\\n",
    "a &= a \\\\\n",
    "b &= r^2 \\exp \\left( b^\\prime \\right)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "The log-transform of parameter $b$ is used in the implementation because\n",
    "taking out $r^2$ causes parameter $b$ to become very small which can\n",
    "cause issues in the optimization process. Note that this also requires the\n",
    "uncertainty of $b^\\prime$ to be transformed to obtain the uncertainty of\n",
    "$b$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pros and cons of both functions\n",
    "\n",
    "There advantages and disadvantages of both implementations are listed below.\n",
    "\n",
    "### Hantush\n",
    "\n",
    "**Pro**:\n",
    "- Parameter A is the gain, which makes it easier to interpret the results.\n",
    "- Estimates the uncertainty of the gain directly.\n",
    "\n",
    "**Con**:\n",
    "- Cannot be used to simulate multiple wells with a single response function.\n",
    "\n",
    "### HantushWellModel\n",
    "\n",
    "**Pro**:\n",
    "- Can be used with WellModel to simulate multiple wells with a single response function.\n",
    "\n",
    "**Con**:\n",
    "- Does not directly estimate the uncertainty of the gain; the uncertainty of the gain must be calculated using special methods.\n",
    "- More sensitive to the initial value of parameters. (The initial parameter values may have to be tweaked to get a good fit result.)\n",
    "\n",
    "\n",
    "So which one should you use? It depends on your use-case:\n",
    "\n",
    "- Use `Hantush` if you are considering a single extraction well or multiple\n",
    "  wells in different aquifers.\n",
    "- Use `HantushWellModel` if you are simulating multiple extraction wells in a\n",
    "  single aquifer or want to pass the distance between an extraction and\n",
    "  observation well as a known parameter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Options\n",
    "\n",
    "Both Hantush implementations in Pastas include two options:\n",
    "\n",
    "- `quad`: numerically integrates the Hantush integrand using `scipy.integrate.quad`. This is relatively slow! (_Note: if available, numba is used to speed up calculation of the integrand_).\n",
    "- `use_numba`: uses [`numba-scipy`](https://github.com/numba/numba-scipy) to speed up calculation of the default Hantush implementation under certain conditions.\n",
    "\n",
    "This yields the following options:\n",
    "- `rf = ps.Hantush()`, default implementation using fast Hantush approximation with numpy (**fast**)\n",
    "- `rf = ps.Hantush(quad=True)`, uses `quad` to numerically integrate Hantush integrand (**slow**)\n",
    "- `rf = ps.Hantush(use_numba=True)`, speeds up fast Hantush approximation with numba-scipy (**fastest**)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The performance is calculated below for the different options listed above\n",
    "(timing may differ depending on your hardware):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rf_numpy = ps.Hantush()\n",
    "rf_quad = ps.Hantush(quad=True)\n",
    "\n",
    "print(\"Hantush approximation (numpy):\")\n",
    "%timeit rf_numpy.step([-0.05, 200, 0.5])\n",
    "print(\"Hantush numerical integration (quad):\")\n",
    "%timeit rf_quad.step([-0.05, 200, 0.5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic example \n",
    "\n",
    "A synthetic example is used to show both Hantush implementations. First, we\n",
    "create a synthetic time series generated with the Hantush response function to\n",
    "which we add autocorrelated residuals. We set the parameter values for the\n",
    "Hantush response function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A defined so that 100 m3/day results in 5 m drawdown\n",
    "Q = 100.0  # m3/day\n",
    "A = -5 / Q\n",
    "a = 200\n",
    "b = 0.5\n",
    "\n",
    "d = 0.0  # reference level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# auto-correlated residuals AR(1)\n",
    "sigma_n = 0.05\n",
    "alpha = 50\n",
    "sigma_r = sigma_n / np.sqrt(1 - np.exp(-2 * 14 / alpha))\n",
    "print(f\"sigma_r = {sigma_r:.2f} m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a head observations time series and a time series with the well extraction rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# head observations between 2000 and 2010\n",
    "idx = pd.date_range(\"2000\", \"2010\", freq=\"D\")\n",
    "ho = pd.Series(index=idx, data=0.0)\n",
    "\n",
    "# extraction of 100 m3/day between 2002 and 2006\n",
    "well = pd.Series(index=idx, data=0.0)\n",
    "well.loc[\"2002\":\"2006\"] = 100.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the synthetic head timeseries based on the extraction rate and the parameters we defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml0 = ps.Model(ho)  # alleen de tijdstippen waarop gemeten is worden gebruikt\n",
    "rm = ps.StressModel(well, ps.Hantush(quad=True), name=\"well\", up=False, settings=\"well\")\n",
    "ml0.add_stressmodel(rm)\n",
    "ml0.set_parameter(\"well_A\", initial=A)\n",
    "ml0.set_parameter(\"well_a\", initial=a)\n",
    "ml0.set_parameter(\"well_b\", initial=b)\n",
    "ml0.set_parameter(\"constant_d\", initial=d)\n",
    "hsynthetic_no_error = ml0.simulate()[ho.index]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_noise = True  # add correlated noise to synthetic head?\n",
    "fit_constant = True  # fit constant separately\n",
    "report = False  # print fit reports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add the auto-correlated residuals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delt = (ho.index[1:] - ho.index[:-1]).values / pd.Timedelta(\"1d\")\n",
    "np.random.seed(1)\n",
    "noise = sigma_n * np.random.randn(len(ho))\n",
    "residuals = np.zeros_like(noise)\n",
    "residuals[0] = noise[0]\n",
    "for i in range(1, len(ho)):\n",
    "    residuals[i] = np.exp(-delt[i - 1] / alpha) * residuals[i - 1] + noise[i]\n",
    "\n",
    "if add_noise:\n",
    "    hsynthetic = hsynthetic_no_error + residuals\n",
    "else:\n",
    "    hsynthetic = hsynthetic_no_error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = hsynthetic_no_error.plot(label=\"synthetic heads (no error)\", figsize=(10, 5))\n",
    "hsynthetic.plot(ax=ax, color=\"C1\", label=\"synthetic heads (with error)\")\n",
    "ax.legend(loc=\"best\")\n",
    "ax.set_ylabel(\"head (m+ref)\")\n",
    "ax.grid(visible=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create three models:\n",
    "\n",
    "1. Model with `Hantush` response function.\n",
    "2. Model with `HantushWellModel` response function with $r$ set to 1.0 m.\n",
    "3. Model with `WellModel`, which uses `HantushWellModel` and $r$ is set to 1.0 m in `WellModel`.\n",
    "\n",
    "All three models should yield the similar results and be able to estimate the true values of the parameters reasonably well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hantush\n",
    "ml_h1 = ps.Model(hsynthetic, name=\"gain\")\n",
    "wm_h1 = ps.StressModel(well, ps.Hantush(), name=\"well\", up=False, settings=\"well\")\n",
    "ml_h1.add_stressmodel(wm_h1)\n",
    "if add_noise:\n",
    "    ml_h1.add_noisemodel(ps.ArNoiseModel())\n",
    "ml_h1.set_parameter(\"constant_d\", initial=0.0)\n",
    "ml_h1.solve(report=report, fit_constant=fit_constant)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve with noise model and HantushWellModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HantushWellModel\n",
    "ml_h2 = ps.Model(hsynthetic, name=\"scaled\")\n",
    "rfunc = ps.HantushWellModel()\n",
    "rfunc.set_distances(1.0)\n",
    "wm_h2 = ps.StressModel(well, rfunc, name=\"well\", up=False, settings=\"well\")\n",
    "ml_h2.add_stressmodel(wm_h2)\n",
    "if add_noise:\n",
    "    ml_h2.add_noisemodel(ps.ArNoiseModel())\n",
    "ml_h2.set_parameter(\"constant_d\", initial=0.0)\n",
    "ml_h2.solve(report=report, fit_constant=fit_constant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WellModel\n",
    "r = np.array([1.0])  # parameter r\n",
    "well.name = \"well\"\n",
    "\n",
    "ml_h3 = ps.Model(hsynthetic, name=\"wellmodel\")\n",
    "wm_h3 = ps.WellModel(\n",
    "    [well], \"well\", r, ps.HantushWellModel(), up=False, settings=\"well\"\n",
    ")\n",
    "ml_h3.add_stressmodel(wm_h3)\n",
    "if add_noise:\n",
    "    ml_h3.add_noisemodel(ps.ArNoiseModel())\n",
    "ml_h3.set_parameter(\"constant_d\", initial=0.0)\n",
    "ml_h3.solve(report=report, fit_constant=fit_constant)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a comparison of all three models. The three models all yield similar results (all the lines overlap). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ps.plots.compare([ml_h1, ml_h2, ml_h3], adjust_height=True, figsize=(10, 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the optimized parameters for each model with the true values we defined at the beginning of this example. Note that we're comparing the value of the gain (not parameter $A$) and that each model has its own method for calculating the gain. As expected, the parameter estimates are reasonably close to the true values defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(\n",
    "    index=[\"well_gain\", \"well_a\", \"well_b\"],\n",
    "    columns=[\"True value\", \"Hantush\", \"HantushWellModel\", \"WellModel\"],\n",
    ")\n",
    "\n",
    "df[\"True value\"] = A, a, b\n",
    "\n",
    "df[\"Hantush\"] = (\n",
    "    # gain (same as A in this case)\n",
    "    wm_h1.rfunc.gain(ml_h1.get_parameters(\"well\")),\n",
    "    # a\n",
    "    ml_h1.parameters.loc[\"well_a\", \"optimal\"],\n",
    "    # b\n",
    "    ml_h1.parameters.loc[\"well_b\", \"optimal\"],\n",
    ")\n",
    "\n",
    "df[\"HantushWellModel\"] = (\n",
    "    # gain (not same as A)\n",
    "    wm_h2.rfunc.gain(ml_h2.get_parameters(\"well\")),\n",
    "    # a\n",
    "    ml_h2.parameters.loc[\"well_a\", \"optimal\"],\n",
    "    # b\n",
    "    np.exp(ml_h2.parameters.loc[\"well_b\", \"optimal\"]),\n",
    ")\n",
    "\n",
    "df[\"WellModel\"] = (\n",
    "    # gain, use WellModel.get_parameters() to get params: A, a, b and r\n",
    "    wm_h3.rfunc.gain(wm_h3.get_parameters(model=ml_h3, istress=0)),\n",
    "    # a\n",
    "    ml_h3.parameters.loc[\"well_a\", \"optimal\"],\n",
    "    # b (multiply parameter value by r^2 for comparison)\n",
    "    np.exp(ml_h3.parameters.loc[\"well_b\", \"optimal\"] * r[0] ** 2),\n",
    ")\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall from earlier that when using `ps.Hantush` the gain and uncertainty of\n",
    "the gain are calculated directly. This is not the case for\n",
    "`ps.HantushWellModel`, so to obtain the uncertainty of the gain when using that\n",
    "response function there is a method called\n",
    "`ps.HantushWellModel.variance_gain()` that computes the variance based on the\n",
    "optimal values and (co)variance of parameters $A'$ and $b'$. There is also a\n",
    "convenience method `ps.WellModel.variance_gain()` that picks up the required\n",
    "variances and covariances from the parent model.\n",
    "\n",
    "The code below shows the calculated gain for each model, and how to calculate\n",
    "the variance and standard deviation of the gain for each model. The results\n",
    "show that the calculated values are all very close, as would be expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataframe\n",
    "var_gain = pd.DataFrame(index=df.columns[1:])\n",
    "\n",
    "# add calculated gain\n",
    "var_gain[\"gain\"] = df.iloc[0, 1:].values\n",
    "\n",
    "# Hantush: variance gain is computed directly\n",
    "var_gain.loc[\"Hantush\", \"var gain\"] = ml_h1.fit.pcov.loc[\"well_A\", \"well_A\"]\n",
    "\n",
    "# HantushWellModel: calculate variance gain explicitly providing values\n",
    "var_gain.loc[\"HantushWellModel\", \"var gain\"] = wm_h2.rfunc.variance_gain(\n",
    "    ml_h2.parameters.loc[\"well_A\", \"optimal\"],  # A\n",
    "    ml_h2.parameters.loc[\"well_b\", \"optimal\"],  # b\n",
    "    ml_h2.fit.pcov.loc[\"well_A\", \"well_A\"],  # var_A\n",
    "    ml_h2.fit.pcov.loc[\"well_b\", \"well_b\"],  # var_b\n",
    "    ml_h2.fit.pcov.loc[\"well_A\", \"well_b\"],  # cov_Ab\n",
    ")\n",
    "\n",
    "# WellModel: calculate variance gain providing only the parent model\n",
    "var_gain.loc[\"WellModel\", \"var gain\"] = wm_h3.variance_gain(ml_h3, istress=0)\n",
    "\n",
    "# calculate std dev gain\n",
    "var_gain[\"std gain\"] = np.sqrt(var_gain[\"var gain\"])\n",
    "\n",
    "# show table\n",
    "var_gain.style.format(\"{:.5e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the true parameters, the estimated parameters and uncertainty ranges \n",
    "($\\pm 2 \\sigma$):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = [ml_h1, ml_h2, ml_h3]\n",
    "\n",
    "fig, axes = plt.subplots(5, 1, sharex=True, figsize=(6, 6))\n",
    "\n",
    "# Gain\n",
    "axes[0].axhline(A, linestyle=\"dashed\", color=\"k\", label=\"True value\")\n",
    "axes[0].errorbar(\n",
    "    range(len(models)),\n",
    "    df.loc[\"well_gain\", \"Hantush\":],\n",
    "    marker=\"o\",\n",
    "    mec=\"k\",\n",
    "    ls=\"none\",\n",
    "    mew=0.5,\n",
    "    yerr=2 * var_gain.loc[:, \"std gain\"],\n",
    "    capsize=3,\n",
    "    ecolor=\"C0\",\n",
    "    label=\"Estimated value\",\n",
    ")\n",
    "axes[0].set_ylabel(\"gain\")\n",
    "\n",
    "# a\n",
    "axes[1].axhline(a, linestyle=\"dashed\", color=\"k\", label=\"True value\")\n",
    "axes[1].errorbar(\n",
    "    range(len(models)),\n",
    "    df.loc[\"well_a\", \"Hantush\":],\n",
    "    marker=\"o\",\n",
    "    mec=\"k\",\n",
    "    ls=\"none\",\n",
    "    mew=0.5,\n",
    "    yerr=2 * np.array([iml.parameters.loc[\"well_a\", \"stderr\"] for iml in models]),\n",
    "    capsize=3,\n",
    "    ecolor=\"C0\",\n",
    "    label=\"Estimated value\",\n",
    ")\n",
    "axes[1].set_ylabel(\"a\")\n",
    "\n",
    "# b (NOTE: transformation of uncertainty necessary for log-transformed parameter)\n",
    "axes[2].axhline(b, linestyle=\"dashed\", color=\"k\", label=\"True value\")\n",
    "# transform log(param) uncertainty to linear uncertainty for HantushWellModel\n",
    "b_stderr = np.array(\n",
    "    [\n",
    "        ml_h1.parameters.loc[\"well_b\", \"stderr\"],\n",
    "        np.exp(ml_h2.parameters.loc[\"well_b\", \"optimal\"])\n",
    "        * ml_h2.parameters.loc[\"well_b\", \"stderr\"],\n",
    "        np.exp(ml_h3.parameters.loc[\"well_b\", \"optimal\"])\n",
    "        * ml_h3.parameters.loc[\"well_b\", \"stderr\"],\n",
    "    ]\n",
    ")\n",
    "axes[2].errorbar(\n",
    "    range(len(models)),\n",
    "    df.loc[\"well_b\", \"Hantush\":],\n",
    "    marker=\"o\",\n",
    "    mec=\"k\",\n",
    "    ls=\"none\",\n",
    "    mew=0.5,\n",
    "    yerr=2 * b_stderr,\n",
    "    capsize=3,\n",
    "    ecolor=\"C0\",\n",
    "    label=\"Estimated value\",\n",
    ")\n",
    "axes[2].set_ylabel(\"b\")\n",
    "\n",
    "# constant_d\n",
    "axes[3].axhline(d, linestyle=\"dashed\", color=\"k\", label=\"True value\")\n",
    "axes[3].errorbar(\n",
    "    range(len(models)),\n",
    "    np.array([iml.parameters.loc[\"constant_d\", \"optimal\"] for iml in models]),\n",
    "    marker=\"o\",\n",
    "    mec=\"k\",\n",
    "    ls=\"none\",\n",
    "    mew=0.5,\n",
    "    yerr=2 * np.array([iml.parameters.loc[\"constant_d\", \"stderr\"] for iml in models]),\n",
    "    capsize=3,\n",
    "    ecolor=\"C0\",\n",
    "    label=\"Estimated value\",\n",
    ")\n",
    "axes[3].set_ylabel(\"constant_d\")\n",
    "\n",
    "# noise_alpha\n",
    "axes[4].axhline(alpha, linestyle=\"dashed\", color=\"k\", label=\"True value\")\n",
    "axes[4].errorbar(\n",
    "    range(len(models)),\n",
    "    np.array([iml.parameters.loc[\"noise_alpha\", \"optimal\"] for iml in models]),\n",
    "    marker=\"o\",\n",
    "    mec=\"k\",\n",
    "    ls=\"none\",\n",
    "    mew=0.5,\n",
    "    yerr=2 * np.array([iml.parameters.loc[\"noise_alpha\", \"stderr\"] for iml in models]),\n",
    "    capsize=3,\n",
    "    ecolor=\"C0\",\n",
    "    label=\"Estimated value\",\n",
    ")\n",
    "axes[4].set_ylabel(\"noise_alpha\")\n",
    "\n",
    "# axes settings\n",
    "axes[-1].set_xticks(range(len(models)))\n",
    "axes[-1].set_xticklabels(df.columns[1:], ha=\"center\")\n",
    "for iax in axes.flat:\n",
    "    iax.grid(True)\n",
    "axes[0].legend(loc=(0, 1), ncol=2, frameon=False)\n",
    "\n",
    "# figure settings\n",
    "fig.align_ylabels()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Checking the calculation of the uncertainty of log-transformed parameter $b$**\n",
    "\n",
    "Relation between $b$ and the transformed parameter $b^\\prime$ is given by:\n",
    "$$\n",
    "b = r^2 \\exp \\left( b^\\prime \\right)\n",
    "$$\n",
    "\n",
    "The formula for transforming the uncertainty (through propagation of uncertainty)\n",
    "$$\n",
    "\\sigma^2_{b} = \\left( \\frac{d b}{d b^\\prime} \\right)^2 \\sigma^2_{b^\\prime}\n",
    "$$\n",
    "\n",
    "The derivative with respect to $b^\\prime$ is obviously:\n",
    "$$\n",
    "\\frac{d b}{d b^\\prime} = r^2 \\exp \\left( b^\\prime \\right)\n",
    "$$\n",
    "Which yields:\n",
    "$$\n",
    "\\sigma^2_b = r^4 \\exp \\left( 2b^\\prime \\right) \\sigma^2_{b^\\prime}\n",
    "$$\n",
    "\n",
    "Testing these formulas on the results obtained with the pastas Models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter b\n",
    "b = ml_h1.parameters.loc[\"well_b\", \"optimal\"]\n",
    "sigma_b = ml_h1.parameters.loc[\"well_b\", \"stderr\"]\n",
    "upper_b = b + 2 * sigma_b\n",
    "lower_b = b - 2 * sigma_b\n",
    "\n",
    "# bprime = log transformed b\n",
    "bp = ml_h2.parameters.loc[\"well_b\", \"optimal\"]\n",
    "sigma_bp = ml_h2.parameters.loc[\"well_b\", \"stderr\"]\n",
    "\n",
    "print(f\"{b=:.4f}, {r[0]**2 * np.exp(bp)=:.4f}\")\n",
    "sigma_b_transformed = r[0] ** 2 * np.exp(bp) * sigma_bp\n",
    "print(f\"{sigma_b=:.4f}, {sigma_b_transformed=:.4f}\")"
   ]
  }
 ],
 "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.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "dace5e1b41a98a8e52d2a8eebc3b981caf2c12e7a76736ebfb89a489e3b62e79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}