{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 07: Using functions to solve an equation-- The Theis example\n",
    "In this exercise we will implement the Theis equation in Python using functions, to evaluate drawdown in hydraulic head from pumping at a well.\n",
    "\n",
    "\n",
    "### Background\n",
    "The Theis (1935) equation is used to calculate drawdown for two-dimensional radial groundwater flow to a point source in an infinite, homogeneous aquifer. The Theis equation was derived from heat transfer literature (with the mathematical help of C.I. Lubin) and is defined as:\n",
    "\n",
    "\\begin{equation}\n",
    "s = \\frac{Q}{4 \\pi T} W(u)\n",
    "\\end{equation}\n",
    "\n",
    "where  \n",
    "$s$ is drawdown [L],  \n",
    "$Q$ is the pumping rate [L$^3$/T],  \n",
    "$T$ is the aquifer transmissivity [L$^2$/T],  \n",
    "$u$ is a dimensionless time parameter [unitless], and  \n",
    "$W(u)$ is the Well function (exponential integral $E_1$) [unitless]. The exponential integral is available in ``scipy.special`` as the ``exp1()`` function.\n",
    "\n",
    "The dimensionless time parameter is defined as:\n",
    "\n",
    "\\begin{equation}\n",
    "u = \\frac{r^2S}{4Tt}\n",
    "\\end{equation}\n",
    "\n",
    "where  \n",
    "$r$ is the distance from the pumping well to a point where drawdown is observed [L],   \n",
    "$S$ is storativity [unitless], and  \n",
    "$t$ is the time since pumping began. \n",
    "\n",
    "Storativity is defined as:\n",
    "\n",
    "\\begin{equation}\n",
    "S = S_s b\n",
    "\\end{equation}\n",
    "\n",
    "where  \n",
    "$S_s$ is specific storage [1/L] and   \n",
    "$b$ is the thickness of the aquifer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.special import exp1 as W"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1)\n",
    "Make a function for the Theis equation. The function should include a [Numpy-style docstring](https://numpydoc.readthedocs.io/en/latest/format.html) that documents the inputs and what the function returns.\n",
    "\n",
    "\\begin{equation}\n",
    "s = \\frac{Q}{4 \\pi T} W(u)\n",
    "\\end{equation}\n",
    "\n",
    "hints: \n",
    "* since this also includes the dimensionless time parameter, the inputs are $r$, $t$, $Q$, $T$, and $S$ \n",
    "* we imported $W$ above from ``scipy.special``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def theis(r, t, Q=1.16, T=100, S=0.0001):\n",
    "    \"\"\"Use the Theis equation to get drawdown\n",
    "    in a confined aquifer at a distance r \n",
    "    from a well pumping at rate Q, for time t.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    r : float or array or floats\n",
    "        Distance to the pumping well (L)\n",
    "    t : float or list-like of floats\n",
    "        Times to calculate drawdown at (T)\n",
    "    Q : float\n",
    "        Pumping rate (L3/T)\n",
    "    T : float\n",
    "        Aquifer transmissivity (L2/T)\n",
    "    S : float\n",
    "        Aquifer storativity\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    s : float array\n",
    "        Drawdown for all values of r or t\n",
    "        \n",
    "    Notes\n",
    "    -----\n",
    "    If r is a 1 or 2D array, t must be scalar.\n",
    "    If t is non-scalar (list-like), r must be scalar.\n",
    "    \n",
    "    Examples\n",
    "    --------\n",
    "    >>> theis(1000, 10, Q=4088, T=1000, S=3e-4)\n",
    "    array([1.40636669])\n",
    "    \"\"\"\n",
    "    \n",
    "    # coerce time input into a numpy array\n",
    "    # (so that we can use the same code \n",
    "    #  for vector or scalar inputs of t)\n",
    "    if np.isscalar(t):\n",
    "        t = [t]\n",
    "    t = np.array(t, dtype=float)\n",
    "    \n",
    "    # dimensionless time parameter\n",
    "    # (only compute for non-zero times)\n",
    "    u = r**2 * S / (4 * T * t)\n",
    "    \n",
    "    s = Q / (4 * np.pi * T) * W(u)\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2)\n",
    "\n",
    "Use the eqn. from Step 1) to compute drawdown at a distance of 1,000m, at 10 days, assuming a pumping rate of 4,088 m/day, transmissivity (T) of 1,000 m/day, and a storativity (S) of 3 x 10$^{-4}$\n",
    "\n",
    "hint: you should get ``1.40636669``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theis(1000, 10, Q=4088, T=1000, S=3e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3)\n",
    "Make a plot of drawdown over time for the same parameters as in Step 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = np.logspace(-1, 1, 100)\n",
    "s = theis(1000, times, Q=4088, T=1000, S=3e-4)\n",
    "\n",
    "plt.plot(times, s)\n",
    "plt.ylabel('Drawdown, in meters')\n",
    "plt.xlabel('Time, in days')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4)\n",
    "\n",
    "\n",
    "Make a new function that accepts arrays of $x$ and $y$ locations, and a pumping well location ($x$, $y$), and returns a corresponding array of $s$ (drawdown) locations, for any number of times. This new function should call the function created in Step 1). To do this, you will need an additional function to compute the distance to the pumping well $r$ for any $x$, $y$ location.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_distance(x0, y0, x1, y1):\n",
    "    \"\"\"Compute the distance between two points.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x0 : float or array-like\n",
    "        x-coordinate(s) for computing distance\n",
    "    y0 : float or array-like\n",
    "        y-coordinate(s) for computing distance\n",
    "    x1 : float\n",
    "        x-coordinate to compute distance from\n",
    "    y1 : float\n",
    "        y-coordinate to compute distance from\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    distance : float or array\n",
    "        Distance of each point defined by x0, y0 from\n",
    "        x1, y1.\n",
    "    \"\"\"\n",
    "    return np.sqrt((x1-x0)**2 + (y1-y0)**2)\n",
    "\n",
    "\n",
    "def theis_xy(x, y, pumping_well_xy, t, Q=1.16, T=100, S=0.0001):\n",
    "    \"\"\"Use the Theis equation to get drawdown\n",
    "    in a confined aquifer at any x, y point(s),\n",
    "    for a pumping well at (x, y) location pumping_well_xy,\n",
    "    pumping at rate Q, for time t.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x : float or list-like of floats\n",
    "        x-coordinates for computing drawdown.\n",
    "    y : float or list-like of floats\n",
    "        y-coordinates for computing drawdown.\n",
    "    pumping_well_xy : tuple\n",
    "        (x, y) location of the pumping well\n",
    "    t : float or list-like of floats\n",
    "        Times to calculate drawdown at (T)\n",
    "    Q : float\n",
    "        Pumping rate (L3/T)\n",
    "    T : float\n",
    "        Aquifer transmissivity (L2/T)\n",
    "    S : float\n",
    "        Aquifer storativity\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    s : list of arrays\n",
    "        Arrays of drawdown (of same shape as x and y),\n",
    "        at each time in t.\n",
    "    \"\"\"\n",
    "    if np.isscalar(t):\n",
    "        t = [t]\n",
    "    if np.isscalar(x):\n",
    "        x = [x]\n",
    "    if np.isscalar(y):\n",
    "        y = [y]\n",
    "    x = np.array(x)\n",
    "    y = np.array(y)\n",
    "    \n",
    "    s = []\n",
    "    for ts in t:\n",
    "        s_t = []\n",
    "        r = get_distance(x, y, *pumping_well_xy)\n",
    "        s_xy = theis(r, t, Q=Q, T=T, S=S)\n",
    "        s_t.append(s_xy)\n",
    "        s_t = np.reshape(s_t, np.shape(x))\n",
    "        s.append(s_t)\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5)\n",
    "\n",
    "Use the function from Step 4) to make a 2D plot of drawdown at 10 days, on a 2,000 x 2,000 meter grid (with an origin of 0, 0), with the pumping well located at 1000, 1000.\n",
    "\n",
    "hint: if the function is taking too long to run, try calling function from step 1) with *vectors* of x,y instead of *scalars*. It should run in less than a second."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "x, y = np.meshgrid(np.arange(2000), np.arange(2000))\n",
    "pumping_well_xy = (1000, 1000)\n",
    "\n",
    "s = theis_xy(x, y, pumping_well_xy, 10, Q=4088, T=1000, S=3e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theis_xy(1000, 0, (1000, 1000), 10, Q=4088, T=1000, S=3e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(s[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6)\n",
    "With $s$ returned as an array for a given set of x, y locations, we can now use *superposition* (adding drawdowns together) to evaluate the effects of multiple pumping wells. Try calling the function from Step 4) again with a well at 500, 500. The $s$ output can be added to the results from Step 5) to get a composite map of drawdown with the two wells pumping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = theis_xy(x, y, (1000, 1000), 10, Q=4088, T=1000, S=3e-4)\n",
    "s2 = theis_xy(x, y, (500, 500), 10, Q=4088, T=1000, S=3e-4)\n",
    "\n",
    "s_total = s1[0] + s2[0]\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "im = ax.imshow(s_total)\n",
    "cs = ax.contour(s_total, colors='w', levels=np.arange(0, 10))\n",
    "ax.clabel(cs, inline=True, fmt='%.2f', fontsize=10)\n",
    "fig.colorbar(im, label='Drawdown, in meters', shrink=0.8, pad=0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}