{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Synthetic Pumping Test - 2 aquifers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from timflow import transient as tft\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (6, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Head data is generated for a pumping test in a two-aquifer model. The well starts pumping at time $t=0$ with a discharge $Q=800$ m$^3$/d. The head is measured in an observation well 10 m from the pumping well. The thickness of the aquifer is 20 m. Questions:\n",
    "\n",
    "1. Determine the optimal values of the hydraulic conductivity and specific storage coefficient of the aquifer when the aquifer is approximated as confined. Use a least squares approach and make use of the `fmin` function of `scipy.optimize` to find the optimal values. Plot the data with dots and the best-fit model in one graph. Print the optimal values of $k$ and $S_s$ to the screen as well as the root mean squared error of the residuals. \n",
    "\n",
    "2. Repeat Question 1 but now approximate the aquifer as semi-confined. Plot the data with dots and the best-fit model in one graph. Print to the screen the optimal values of $k$, $S_s$ and $c$  to the screen as well as the root mean squared error of the residuals. Is the semi-cofined model a better fit than the confined model?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data():\n",
    "    # 2 layer model with some random error\n",
    "    ml = tft.ModelMaq(\n",
    "        kaq=[10, 20],\n",
    "        z=[0, -20, -22, -42],\n",
    "        c=[1000],\n",
    "        Saq=[0.0002, 0.0001],\n",
    "        tmin=0.001,\n",
    "        tmax=100,\n",
    "    )\n",
    "    tft.Well(ml, 0, 0, rw=0.3, tsandQ=[(0, 800)])\n",
    "    ml.solve()\n",
    "    t = np.logspace(-2, 1, 100)\n",
    "    h = ml.head(10, 0, t)\n",
    "    plt.figure()\n",
    "    r = 0.01 * rnd.random(100)\n",
    "    n = np.zeros_like(r)\n",
    "    # alpha = 0.8\n",
    "    for i in range(1, len(n)):\n",
    "        n[i] = 0.8 * n[i - 1] + r[i]\n",
    "    ho = h[0] + n\n",
    "    plt.plot(t, ho, \".\")\n",
    "    data = np.zeros((len(ho), 2))\n",
    "    data[:, 0] = t\n",
    "    data[:, 1] = ho\n",
    "    # np.savetxt('pumpingtestdata.txt', data, fmt='%2.3f', header='time (d), head (m)')\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnd = np.random.default_rng(11)\n",
    "data = generate_data()\n",
    "to = data[:, 0]\n",
    "ho = data[:, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(p, to=to, ho=ho, returnmodel=False):\n",
    "    k = p[0]\n",
    "    S = p[1]\n",
    "    ml = tft.ModelMaq(kaq=k, z=[0, -20], Saq=S, tmin=0.001, tmax=100)\n",
    "    tft.Well(ml, 0, 0, rw=0.3, tsandQ=[(0, 800)])\n",
    "    ml.solve(silent=True)\n",
    "    if returnmodel:\n",
    "        return ml\n",
    "    h = ml.head(10, 0, to)\n",
    "    return np.sum((h[0] - ho) ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import fmin\n",
    "\n",
    "lsopt = fmin(func, [10, 1e-4])\n",
    "print(\"optimal parameters:\", lsopt)\n",
    "print(\"rmse:\", np.sqrt(func(lsopt) / len(ho)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = func(lsopt, returnmodel=True)\n",
    "plt.figure()\n",
    "plt.plot(data[:, 0], data[:, 1], \".\", label=\"observed\")\n",
    "hm = ml.head(10, 0, to)\n",
    "plt.plot(to, hm[0], \"r\", label=\"modeled\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time (d)\")\n",
    "plt.ylabel(\"head (m)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cal = tft.Calibrate(ml)\n",
    "cal.set_parameter(name=\"kaq\", layers=0, initial=10, pmin=0.1, pmax=1000)\n",
    "cal.set_parameter(name=\"Saq\", layers=0, initial=1e-4, pmin=1e-5, pmax=1e-3)\n",
    "cal.series(name=\"obs1\", x=10, y=0, layer=0, t=to, h=ho)\n",
    "cal.fit(report=False)\n",
    "print(\"rmse:\", cal.rmse())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cal.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model as semi-confined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func2(p, to=to, ho=ho, returnmodel=False):\n",
    "    k = p[0]\n",
    "    S = p[1]\n",
    "    c = p[2]\n",
    "    ml = tft.ModelMaq(\n",
    "        kaq=k, z=[2, 0, -20], Saq=S, c=c, topboundary=\"semi\", tmin=0.001, tmax=100\n",
    "    )\n",
    "    tft.Well(ml, 0, 0, rw=0.3, tsandQ=[(0, 800)])\n",
    "    ml.solve(silent=True)\n",
    "    if returnmodel:\n",
    "        return ml\n",
    "    h = ml.head(10, 0, to)\n",
    "    return np.sum((h[0] - ho) ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lsopt2 = fmin(func2, [10, 1e-4, 1000])\n",
    "print(\"optimal parameters:\", lsopt2)\n",
    "print(\"rmse:\", np.sqrt(func2(lsopt2) / len(ho)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = func2(lsopt2, returnmodel=True)\n",
    "plt.figure()\n",
    "plt.plot(data[:, 0], data[:, 1], \".\", label=\"observed\")\n",
    "hm = ml.head(10, 0, to)\n",
    "plt.plot(to, hm[0], \"r\", label=\"modeled\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time (d)\")\n",
    "plt.ylabel(\"head (m)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tft.ModelMaq(\n",
    "    kaq=10, z=[2, 0, -20], Saq=1e-4, c=1000, topboundary=\"semi\", tmin=0.001, tmax=100\n",
    ")\n",
    "w = tft.Well(ml, 0, 0, rw=0.3, tsandQ=[(0, 800)])\n",
    "ml.solve(silent=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cal = tft.Calibrate(ml)\n",
    "cal.set_parameter(name=\"kaq\", layers=0, initial=10)\n",
    "cal.set_parameter(name=\"Saq\", layers=0, initial=1e-4)\n",
    "cal.set_parameter(name=\"c\", layers=0, initial=1000)\n",
    "cal.series(name=\"obs1\", x=10, y=0, layer=0, t=to, h=ho)\n",
    "cal.fit(report=False)\n",
    "cal.parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cal.rmse(), ml.aq.kaq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(data[:, 0], data[:, 1], \".\", label=\"observed\")\n",
    "hm = ml.head(10, 0, to)\n",
    "plt.plot(to, hm[0], \"r\", label=\"modeled\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time (d)\")\n",
    "plt.ylabel(\"head (m)\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}