gw_modelling_python_session_03_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f4ee7a8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Groundwater Modelling in Python - Session 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c262a0a4-6132-443a-a472-243be6908bbb",
   "metadata": {},
   "source": [
    "<img src=\"figs/part_of_cover_bakker_post.png\" width=\"800px\"> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e62b54a2-dcb3-4512-8fc9-85d871e7e2a9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Before starting we need to import NumPy and Matplotlib and set some defaults\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# and set some parameters to make the figures look good\n",
    "plt.rcParams[\"figure.autolayout\"] = True # same as tight_layout after every plot\n",
    "plt.rcParams[\"figure.figsize\"] = (9, 4) # set default figure size\n",
    "plt.rcParams[\"contour.negative_linestyle\"] = 'solid' # set default line style\n",
    "plt.rc('font', size=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1361e7de-488b-40bb-aeef-85739f75ff4f",
   "metadata": {},
   "source": [
    "Sections and Exercises with a star are additional material that is not covered in detail during the workshop."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d395dcf5-61ce-4283-8bc1-4aaaebebaf8b",
   "metadata": {},
   "source": [
    "# Introduction to two-dimensional flow\n",
    "The discharge vector for two-dimensional flow in a horizontal plane consists of two components: $Q_x$ in the $x$-direction and $Q_y$ in the $y$-direction. Continuity of flow for steady two-dimensional flow in the horizontal $x, y$ plane may be written as \n",
    "\\begin{equation}\\label{basicdeq2d}\n",
    "\\frac{\\partial Q_x}{\\partial x} + \\frac{\\partial Q_y}{\\partial y} = N\n",
    "\\end{equation}\n",
    "where $N$ is the areal recharge. \n",
    "Note that this equation differs from the continuity of flow equation for one-dimensional steady flow only by the additional term $\\partial Q_y / \\partial y$. \n",
    "\n",
    "Steady confined and unconfined flow may be formulated as potential flow, where the discharge vector is minus the gradient of the discharge potential\n",
    "\\begin{equation}\\label{QxQychap6}\n",
    "Q_x = -\\frac{\\partial \\Phi}{\\partial x} \\quad\\quad Q_y = -\\frac{\\partial \\Phi}{\\partial y}\n",
    "\\end{equation}\n",
    "so that the differential equation for steady potential flow becomes the two-dimensional Poisson equation\n",
    "\\begin{equation}\\label{poisson2d}\n",
    "\\nabla^2 \\Phi = \\frac{\\partial^2 \\Phi}{\\partial x^2} + \\frac{\\partial^2 \\Phi}{\\partial y^2} = -N\n",
    "\\end{equation}\n",
    "In absence of areal recharge ($N=0$), this differential equation is referred to as the two-dimensional Laplace equation.\n",
    "\n",
    "The relationship between the discharge potential and head for the different flow types is summarized as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97a7c7d4-2ffc-40c8-a3d8-4f70d90babda",
   "metadata": {},
   "source": [
    "\\begin{equation}\\label{potential_table}\n",
    "\\begin{array}{|c|cc|}\n",
    "\\hline\n",
    "\\text{Flow type} & \\text{Discharge potential} & \\text{Condition}\\\\\n",
    "\\hline\n",
    "\\text{confined} & \\Phi = kHh & h \\ge z_t \\\\\n",
    "\\hline\n",
    "\\text{unconfined} & \\Phi = \\tfrac{1}{2}k(h-z_b)^2 & h \\ge z_b \\\\\n",
    "\\hline\n",
    "\\text{unconfined/confined} & \\Phi = kHh -(\\tfrac{1}{2}kH^2 + kHz_b) & h \\ge z_t \\\\\n",
    "\\text{flow} & \\Phi = \\tfrac{1}{2}k(h-z_b)^2 & z_b \\le h \\le z_t \\\\\n",
    "\\hline\n",
    "\\end{array}\n",
    "\\end{equation}\n",
    "\n",
    "where $k$ is the hydraulic conductivity of the aquifer, $z_b$ and $z_t$ are the elevations of the bottom and top of the aquifer, and $H$ is the aquifer thickness."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3c9388b-fca5-40df-b1f7-ac677eaab185",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Solution 1. Flow to a pumping well on an island with infiltration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080c51b8-ec4c-435d-aea8-8927b72a91f6",
   "metadata": {},
   "source": [
    "Consider a pumping well with discharge $Q$ at the center of a circular island of radius $R$. The aquifer is unconfined with hydraulic conductivity $k$ and base elevation $z_b$. The water level is fixed at the edge of the island to $h_R$. The recharge on the island is uniform and equal to $N$. Note that the discharge $Q$ is positive when the well is pumping water out of the aquifer and the recharge $N$ is positive when water is entering the aquifer. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c76723fe-64d6-434c-b227-fe8ae212b9bd",
   "metadata": {},
   "source": [
    "<img src=\"figs/merged_nbs6_17_0.png\" width=\"400px\"> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea36ee44-17f1-4bed-bba0-7d2e0391dcd4",
   "metadata": {},
   "source": [
    "The solution for the discharge potential is\n",
    "\\begin{equation} \\label{phi3}\n",
    "\\Phi = -\\frac{1}{4}N (r^2-R^2) + \\frac{Q}{4\\pi} \\ln{(r^2/R^2)}\n",
    "+\\Phi_R\n",
    "\\end{equation}\n",
    "where $\\Phi_R=\\tfrac{1}{2} k  (h_R - z_b)^2$\n",
    "\n",
    "The head in the aquifer is computed from the discharge potential as\n",
    "\\begin{equation}\n",
    "h = z_b + \\sqrt{\\frac{2\\Phi}{k}}\n",
    "\\end{equation}\n",
    "\n",
    "As a first example, the discharge of the well is half the total recharge on the island \n",
    "$Q = N \\pi R^2 / 2$. The head in the well is equal to the head at $r=r_w$, where $r_w$ is the radius of the well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b805f50e-7c5f-4ead-be04-99343dbe3c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "k = 5 # hydraulic conductivity, m/d\n",
    "zb = 0 # bottom elevation of aquifer, m\n",
    "N = 0.001 # areal recharge, m/d\n",
    "R = 200 # radius of the island, m\n",
    "hR = 10  # head at edge of island, m\n",
    "rw = 0.3 # well radius, m\n",
    "Q = 0.5 * N * np.pi * R ** 2  # discharge is half total recharge, m^3/d\n",
    "print(f'Discharge of well: {Q:.2f} m^3/d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b2113d2-e6fc-40d7-9fcb-8d60717369f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solution\n",
    "r = np.linspace(rw, R, 100)\n",
    "phiR = 0.5 * k * (hR - zb) ** 2\n",
    "phi = -N / 4 * (r ** 2 - R ** 2) + Q / (4 * np.pi) * np.log(r ** 2 / R ** 2) + phiR\n",
    "h = zb + np.sqrt(2 * phi / k)\n",
    "# plot\n",
    "plt.plot(r, h, 'C0')\n",
    "plt.plot(-r, h, 'C0')\n",
    "plt.xlabel('radial distance (m)')\n",
    "plt.ylabel('head (m)')\n",
    "plt.xticks(np.arange(-200, 201, 50), np.abs(np.arange(-200, 201, 50))) # positive r ticks\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0346a30-f384-499a-8867-0a9d250dc5e9",
   "metadata": {},
   "source": [
    "## Exercise 1.1\n",
    "Compute the head in the well when the discharge is equal to half the total recharge on the island and when the discharge is equal to exactly the total recharge on the island. Plot the head vs. $r$ for both cases in one graph. Also include the case that $Q=0$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f5b4e71-cef7-44a3-81ea-78730b85d569",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d83bb96c-f397-4208-9d0e-0ceb6a1e511d",
   "metadata": {},
   "source": [
    "## Exercise 1.2$^*$\n",
    "What is the (theoretically) maximum discharge of the well? What happens when you try to compute the head in the well when the discharge is larger than the maximum discharge?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ce8c3c2-0ee9-41ac-91fd-8fe246ac4bbe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec291e5-da32-4275-a44c-ebc3f08ca6ce",
   "metadata": {},
   "source": [
    "# Solution 2. A pumping well near a river"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a07fdb9-2b36-48a4-a650-f9b8ceb9c0e5",
   "metadata": {},
   "source": [
    "A well is located near a long and straight river running North-South along the $y$-axis. The well is located on the left side of the river at $(x_w,y_w)$ and has a discharge $Q$ (see Figure). The river is in full hydraulic contact with the aquifer so that the head in the aquifer is equal to the river stage $h_0$ along the river. The aquifer is unconfined, but the variation of the saturated thickness is neglected: the transmissivity of the aquifer is approximated as constant and equal to $T$. Far away from the well, the head in the aquifer is equal to $h_0$. \n",
    "\n",
    "<img src=\"figs/wellriver.png\" width=\"350px\"> \n",
    "\n",
    "The solution for the head and flow in the aquifer are obtained with the method of images as:\n",
    "\n",
    "\\begin{equation}\n",
    "h = \\frac{Q}{4\\pi T}\\ln\\left[\\frac{(x-x_w)^2 + (y-y_w)^2}{(x+x_w)^2 + (y-y_w)^2}\\right] + h_0\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc60d3d-d87f-4b39-87d8-162abc9bb944",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "Q_x = -T\\frac{\\partial h}{\\partial x} = -\\frac{Q}{2\\pi}\\left[\\frac{x-x_w}{(x-x_w)^2 + (y-y_w)^2} - \n",
    "\\frac{x+x_w}{(x+x_w)^2 + (y-y_w)^2} \\right]\n",
    "\\end{equation}\n",
    "\\begin{equation}\n",
    "Q_y = -T\\frac{\\partial h}{\\partial y} = -\\frac{Q}{2\\pi}\\left[\\frac{y-y_w}{(x-x_w)^2 + (y-y_w)^2} - \n",
    "\\frac{y-y_w}{(x+x_w)^2 + (y-y_w)^2} \\right]\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fdf6bcf-1263-436c-93a6-5cb07eeeadb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def head(x, y, xw, yw, Q, T, h0):\n",
    "    h = Q / (4 * np.pi * T) * np.log(\n",
    "    ((x - xw) ** 2 + (y - yw) ** 2) /\n",
    "    ((x + xw) ** 2 + (y - yw) ** 2)) + h0\n",
    "    return h\n",
    "\n",
    "def QxQy(x, y, xw, yw, Q):\n",
    "    r1sq = (x - xw) ** 2 + (y - yw) ** 2\n",
    "    r2sq = (x + xw) ** 2 + (y - yw) ** 2\n",
    "    Qx = -Q / (2 * np.pi) * ((x - xw) / r1sq - (x + xw) / r2sq)\n",
    "    Qy = -Q / (2 * np.pi) * ((y - yw) / r1sq - (y - yw) / r2sq)\n",
    "    return np.array([Qx, Qy])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecc79f63-f095-470b-b41b-cee7864b48ab",
   "metadata": {},
   "source": [
    "An example is provided below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "417b7436-f15f-4217-ae90-59c872ff4b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "k = 10 # hydraulic conductivity, m/d\n",
    "H = 20 # m\n",
    "xw = -50 # x-location of well, m\n",
    "yw = 0 # y-location of well, m\n",
    "Q = 800 # discharge of well, m^3/d\n",
    "rw = 0.3 # radius of well, m\n",
    "h0 = 20 # river level, m\n",
    "n = 0.3 # porosity, -\n",
    "T = k * H # transmissivity, m^2/d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2590b18-9bbe-4b7c-90bf-9f4e77098831",
   "metadata": {},
   "source": [
    "The head is computed on a grid of 100 by 100 points where $x$ varies from $-100$ to $0$ and $y$ varies from $-50$ to $+50$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "541994fe-c1ae-46b5-80a4-017bf813072b",
   "metadata": {},
   "outputs": [],
   "source": [
    "xg, yg = np.meshgrid(np.linspace(-100, 0, 100), np.linspace(-50, 50, 100))\n",
    "h = head(xg, yg, xw, yw, Q, T, h0)\n",
    "print(f'min and max head in model window: {h.min():.1f} m, {h.max():.1f} m')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17686bb6-3ef3-41d9-bb21-d695133a43c0",
   "metadata": {},
   "source": [
    "The head is contoured as follows. First, a figure is created with the `subplot` command. The first argument is `121`, which means a figure with 1 row of graphs, 2 colulmns of graphs, and start with graph 1 (the left graph; the rigth graph will be created in the next code block). The keyword `aspect=1` is important, as it makes the scale along the $x$ and $y$ axes equal, so that a circle looks like a circle rather than an ellipse.  The `contour` function returns a contour set, which is stored in the variable `cs`. Labels are placed along the contours by passing `cs` to the `clabel` function. The format of the labels is defined ast `'%1.1f'`, wich stands for a float with 1 decimal place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ba6fd4-ecbe-49d5-87e9-35adf5fa61e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, levels=np.arange(15, 21, 0.2), colors='C0')\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot(xw, yw, 'k.') # black dot for  well\n",
    "plt.axvline(0, color='skyblue', lw=8); # blue line for river"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a047db8-e3a1-4b03-9b64-3a55cc80ce0d",
   "metadata": {},
   "source": [
    "The flow field is visualized using the `streamplot` function. First, $Q_x$ and $Q_y$ are computed on the same grid as used for the head. Then the `streamplot` function is used; the `density` keyword is set to control the density of streamlines. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b7d7740-dc12-4a7d-a360-ae1c2e474d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "Qx, Qy = QxQy(xg, yg, xw, yw, Q)\n",
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(15, 21, 0.2), colors='C0')\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot(xw, yw, 'k.') # black dot at location of well\n",
    "plt.axvline(0, color='skyblue', lw=8); # blue line for river\n",
    "#\n",
    "plt.subplot(122, aspect=1)\n",
    "plt.streamplot(xg, yg, Qx, Qy, color='C1', density=0.9)\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2073d95f-4981-4720-a061-09c592541f7b",
   "metadata": {},
   "source": [
    "Alternatively, the `streamplot` function may be used to start a number of streamlines from specified starting points. In the code cell below, 20 starting points are defined, equally distributed along the circumference of the well. Note that streamlines are integrated in both directions from the starting points (with the flow and against the flow) by default. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "554a8239-8456-40ee-bb1f-970daee6bcc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "xstart = xw + rw * np.cos(np.arange(0, 2 * np.pi, 2 * np.pi / 20))\n",
    "ystart = yw + rw * np.sin(np.arange(0, 2 * np.pi, 2 * np.pi / 20))\n",
    "plt.subplot(111, aspect=1)\n",
    "for i in range(20):\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xstart[i], ystart[i])])\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d93122ad-92c0-4902-a483-52ffadc726ec",
   "metadata": {},
   "source": [
    "## Travel time from the river to the well\n",
    "The travel time may be computed through integration of the velocity vector along a streamline. This is somewhat complicated, as most strealines are curved. It is straighforward, however, for the straight streamline starting at $(x,y)=(0,0)$ (recall that $y_w=0$) on the river and flowing towards the well. This travel time represents the shortest travel time from the river to the well, which is important to assess the vulnerability of the well to possible contamination by pathogens and harmful chemicals in the river water.  \n",
    "\n",
    "\\begin{equation}\n",
    "v_x = \\frac{Q_x}{nH} = \\frac{\\text{d}x}{\\text{d}t}\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "t = \\int_{a}^{b}\\frac{\\text{d}x}{v_x}\n",
    "\\end{equation}\n",
    "\n",
    "The integral may computed numerically, for example use the `quad` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc85e562-cc2a-4fca-99fc-f4b2589636e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrand(x, y, xw, yw, Q, H=10, n=0.3):\n",
    "    Qx, Qy = QxQy(x, y, xw, yw, Q)\n",
    "    vx = Qx / (n * H)\n",
    "    return 1 / vx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe1e8ba-5b63-415e-b290-68a8b97aeddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.integrate import quad\n",
    "traveltime = quad(integrand, a=0, b=xw + rw, args=(yw, xw, yw, Q))[0]\n",
    "print(f'shortest travel time from river to well: {traveltime:.1f} d')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f514b9-4f94-402d-8ef8-7f3c140915ef",
   "metadata": {},
   "source": [
    "## Exercise 2.1\n",
    "Move the well farther away from the river such that the shortest travel time from the river to the well is at least 40 days. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "543323fa-b8b9-4e31-8854-29d3cf6d30e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a3f93a0-261c-4524-8a71-d6a509fe6143",
   "metadata": {},
   "source": [
    "## Exercise 2.2$^*$\n",
    "Compute the head at the well at the new location (computed in Exercise 2.1) and compare it to the head at the well when it was located at the original location ($x_w=-50$). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c215153-2823-4564-b2d9-9f9b4f1a8e1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "279a9639-09cc-4e54-841a-ecce5fa10bef",
   "metadata": {},
   "source": [
    "## Multiple wells\n",
    "The solution for multiple wells may be obtained through superposition. For example, consider two wells located on the western side of a long and straight river running North-South along the $y$-axis. The first well has discharge $Q_0$ and is located at $(x_0,y_0)$. The second well has discharge $Q_1$ and is located at $(x_1,y_1)$. The solution is\n",
    "\\begin{equation}\n",
    "h = \\frac{Q_0}{4\\pi T}\\ln\\left[\\frac{(x-x_0)^2 + (y-y_0)^2}{(x+x_0)^2 + (y-y_0)^2}\\right] + \\frac{Q_1}{4\\pi T}\\ln\\left[\\frac{(x-x_1)^2 + (y-y_1)^2}{(x+x_1)^2 + (y-y_1)^2}\\right] + h_0\n",
    "\\end{equation}\n",
    "The radius of both wells is $r_w$. \n",
    "The head is contoured and a streamplot is created in the code cell below for a doublet, an extraction and injection well with equal but opposite discharge (i.e., $Q_1=-Q_0$). Use is made of the `head` function defined previously. Note that part of the injected water flows to the river and part of the extracted water comes from the river. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a0436d-38b8-4e6d-a85d-9457e40d467c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# additional parameters\n",
    "Q0 = 800 # discharge of well, m^3/d\n",
    "x0 = -100 # m\n",
    "y0 = 50 # m\n",
    "Q1 = -Q0\n",
    "x1 = -100 # m\n",
    "y1 = -50 # m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbbe792a-3e3c-4279-a23a-b667db87d829",
   "metadata": {},
   "outputs": [],
   "source": [
    "xg, yg = np.meshgrid(np.linspace(-200, 0, 100), np.linspace(-100, 100, 100))\n",
    "h = head(xg, yg, x0, y0, Q0, T, 0) + head(xg, yg, x1, y1, Q1, T, 0) + h0\n",
    "print(f'min and max head in model window: {h.min():.1f}, {h.max():.1f} m')\n",
    "Qx, Qy = QxQy(xg, yg, x0, y0, Q0) + QxQy(xg, yg, x1, y1, Q1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "408b89a3-6176-4006-a060-08684777d413",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(17, 23, 0.2), colors='C0')\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot([x0, x1], [y0, y1], 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8)\n",
    "#\n",
    "plt.subplot(122, aspect=1)\n",
    "plt.streamplot(xg, yg, Qx, Qy, color='C1', density=0.9) # density of streamlines\n",
    "plt.plot([x0, x1], [y0, y1], 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e304b38c-2e62-45bb-997f-145c64b22aa1",
   "metadata": {},
   "source": [
    "## Exercise 2.3$^*$\n",
    "Compute the head at the extraction well and the head at the injection well for the above case of a doublet near a river."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca0868ab-7477-4b67-9583-c12815c68c42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38ab6a6-8e98-475e-b6df-c83c44337f8a",
   "metadata": {},
   "source": [
    "# Solution 3. A pumping well near a river in uniform background flow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc22112a-d269-4121-9cb4-4cf3ea1b4526",
   "metadata": {},
   "source": [
    "In virtually all aquifers, the natural groundwater flow is towards a stream, river, lake, or ocean. This flow is also called the background flow, as it is present prior to the installation of a well. The background flow may be estimated, e.g., from the heads in two observation wells. For example, the flow in the $x$-direction may be obtained from two observation wells aligned in the $x$-direction as\n",
    "\\begin{equation}\n",
    "U\\approx T\\frac{\\Delta h}{\\Delta x}\n",
    "\\end{equation}\n",
    "where $\\Delta h$ is the head drop between the observation wells and $\\Delta x$ is the distance between the two observation wells. \n",
    "\n",
    "Here, we are considering a pumping well near a long and straight river running North-South (as in the previous solution), but now with a uniform background flow towards the river (i.e., in the $x$-direction), as shown in the figure.\n",
    "\n",
    "<img src=\"figs/wellriveruflow.png\" width=\"350px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d8e0f6-bb67-4cf8-b816-dfc7c3730b21",
   "metadata": {},
   "source": [
    "The solution for the head in the aquifer is obtained from the previous solution by adding a uniform flow term in the $x$-direction. The solutions for the head and flow are\n",
    "\n",
    "\\begin{equation}\n",
    "h = \\frac{Q}{4\\pi T}\\ln\\left[\\frac{(x-x_w)^2 + (y-y_w)^2}{(x+x_w)^2 + (y-y_w)^2}\\right] \n",
    "-\\frac{U}{T}x + h_0\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "Q_x = -\\frac{Q}{2\\pi}\\left[\\frac{x-x_w}{(x-x_w)^2 + (y-y_w)^2} - \n",
    "\\frac{x+x_w}{(x+x_w)^2 + (y-y_w)^2} \\right] + U\n",
    "\\end{equation}\n",
    "\\begin{equation}\n",
    "Q_y = -\\frac{Q}{2\\pi}\\left[\\frac{y-y_w}{(x-x_w)^2 + (y-y_w)^2} - \n",
    "\\frac{y-y_w}{(x+x_w)^2 + (y-y_w)^2} \\right]\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37e47603-8f64-429a-81c2-ba573b2da592",
   "metadata": {},
   "source": [
    "Consider a well in a uniform background flow near a river. The well is a bit further from the river and has a somewhat smaller discharge than in the previous example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57524c47-794f-4c45-8235-fcec86cdc6ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "k = 20 # hydraulic conductivity, m/d\n",
    "H = 10 # m\n",
    "xw = -200 # x-location of well, m\n",
    "yw = 0 # y-location of well, m\n",
    "Q = 200 # discharge of well, m^3/d\n",
    "rw = 0.3 # radius of well, m\n",
    "h0 = 20 # river level, m\n",
    "n = 0.3 # porosity, -\n",
    "T = k * H # transmissivity, m^2/d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f07a52-c1c2-4b8b-aefc-693361c4a51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters uniform background flow\n",
    "delx = 1000 # distance between observation wells\n",
    "delh = 2 # head drop between observation wells\n",
    "U = T * delh / delx\n",
    "print(f'Uniform flow U = {U:.1f} m^2/d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae13b60-fdb6-491b-962f-25f6ba57a190",
   "metadata": {},
   "outputs": [],
   "source": [
    "def head(x, y, xw, yw, Q, T, h0, U=0):\n",
    "    h = Q / (4 * np.pi * T) * np.log(\n",
    "    ((x - xw) ** 2 + (y - yw) ** 2) /\n",
    "    ((x + xw) ** 2 + (y - yw) ** 2)) -U * x / T + h0\n",
    "    return h\n",
    "\n",
    "def QxQy(x, y, xw, yw, Q, U=0):\n",
    "    r1sq = (x - xw) ** 2 + (y - yw) ** 2\n",
    "    r2sq = (x + xw) ** 2 + (y - yw) ** 2\n",
    "    Qx = -Q / (2 * np.pi) * ((x - xw) / r1sq - (x + xw) / r2sq) + U\n",
    "    Qy = -Q / (2 * np.pi) * ((y - yw) / r1sq - (y - yw) / r2sq)\n",
    "    return Qx, Qy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48241dbc-e1cc-466d-8564-7e98ec7dad95",
   "metadata": {},
   "source": [
    "The head is contoured and a streamplot is created. 20 streamlines are started from $x=-600$ m, equally distribute between $y=-300$ and $y=+300$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1863b61b-d97a-40e7-b595-3fa18119d0bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "xg, yg = np.meshgrid(np.linspace(-600, 0, 100), np.linspace(-300, 300, 100))\n",
    "h = head(xg, yg, xw, yw, Q, T, h0, U)\n",
    "print(f'min and max head in window: {h.min():.1f} m, {h.max():.1f} m')\n",
    "Qx, Qy = QxQy(xg, yg, xw, yw, Q, U)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ca8cfb-c4df-4738-aec9-677073b402b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(19, 22, 0.1), colors='C0')\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8)\n",
    "#\n",
    "plt.subplot(122, aspect=1)\n",
    "xstart = -600 * np.ones(20)\n",
    "ystart = np.linspace(-300, 300, 20)\n",
    "for i in range(20):\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xstart[i], ystart[i])], \n",
    "                   integration_direction='forward')\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47e0b69f-71ce-4377-b647-0fbafe22dbee",
   "metadata": {},
   "source": [
    "There appears to be a stagnation point (a point with zero velocity) between the well and the river. This point may be found by setting $Q_x=0$ along the line $y=y_w$ and solving for $x$, which gives the location $x_s$ of the stagnation point:\n",
    "\\begin{equation}\n",
    "x_s = -\\sqrt{x_w^2 + \\frac{Qx_w}{\\pi U}}\n",
    "\\end{equation}\n",
    "Two streamlines are started from the stagnation point, one just above the stagnation point and one just below the stagnation point. The streamlines are integrated `backward`, so against the flow. These two streamlines bound the capture zone envelope, the area of the aquifer that is (eventually) captured by the well. Note that far upstream the width of the capture zone envelope is $W=Q/U$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67e9cac6-8c98-4a33-9c72-d8679d883790",
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = -np.sqrt(xw ** 2 + Q * xw / (np.pi * U))\n",
    "plt.subplot(121, aspect=1)\n",
    "plt.contour(xg, yg, h, np.arange(19, 22, 0.1), colors='C0')\n",
    "for ystart in [0.1, -0.1]:\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xs, ystart)], \n",
    "                   integration_direction='backward', color='k')\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8)\n",
    "print(f'width of capture zone envelope far upstream: {Q / U:.1f} m')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c73bda1-db2a-421f-b73f-701a02f21514",
   "metadata": {},
   "source": [
    "### Exercise 3.1\n",
    "Consider a well located in a uniform flow field near a river running North-South, as in the previous example. Near the well is a contaminated area as shown in the figure below (run the code cell). For the current discharge of the well ($Q=200$ m$^3$/d), the well captures water from the contaminated area, as can be seen by starting a streampath from the corner of the contaminated area. The discharge of the well must be decreased such that it doesn't capture water from the contaminated area. Determine the maximum discharge of the well such that the well doesn't capture water from the contaminated area."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c3d4963-e093-4b9d-8484-feb9d24c83b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run this code cell to see the contaminated area\n",
    "Q = 200\n",
    "Qx, Qy = QxQy(xg, yg, xw, yw, Q, U)\n",
    "plt.subplot(121, aspect=1)\n",
    "plt.plot(xw, yw, 'k.')\n",
    "plt.axvline(0, color='skyblue', lw=8)\n",
    "plt.fill([-300, -100, -100, -300], [100, 100, 300, 300], color='pink')\n",
    "plt.text(-200, 200, 'contaminated', ha='center', va='center', rotation=-45)\n",
    "plt.xlim(-600, 0)\n",
    "plt.ylim(-300, 300)\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5fb6ecd-8fb4-49be-806b-7b518a2afb4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5933893-3ebe-4236-af56-79265b5a2134",
   "metadata": {},
   "source": [
    "### Maximum discharge before capturing river water\n",
    "When the discharge of the well is large, there is no stagnation point between the well and the river. This can be seen from the equation for the stagnation point, as the term under the square root sign becomes negative (recall that $x_w$ is negative). The maximum discharge of the well such that it extracts no river water is when the stagnation point is exactly on the river and $x_s=0$. This gives for the maximum discharge \n",
    "\\begin{equation}\n",
    "Q_\\text{max}=-\\pi U x_w\n",
    "\\end{equation}\n",
    "The maximum discharge is computed below and the capture zone envelope is plotted. Note that the capture zone envelope touches the river. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f5b416-1355-4826-bc00-e0f28b8d50c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "Qmax = -np.pi * U * xw\n",
    "print(f'maximum discharge: {Qmax:.1f} m^3/d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c488b04c-36aa-421c-9efe-d1568454f61d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute grid of Qx, Qy\n",
    "xg, yg = np.meshgrid(np.linspace(-600, 0, 100), np.linspace(-300, 300, 100))\n",
    "Qx, Qy = QxQy(xg, yg, xw, yw, Qmax, U)\n",
    "# plot streamlines\n",
    "plt.subplot(121, aspect=1)\n",
    "plt.axvline(0, color='skyblue', lw=8) # river\n",
    "xstart = -600 * np.ones(20)\n",
    "ystart = np.linspace(-300, 300, 20)\n",
    "for i in range(20):\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xstart[i], ystart[i])], \n",
    "                   integration_direction='forward')\n",
    "# capture zone envelope\n",
    "for y in [-0.1, 0.1]:\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(0, y)], \n",
    "                   integration_direction='backward', color='k')\n",
    "plt.plot(xw, yw, 'k.');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65975f48-11eb-4a6d-913e-72a3a77c9fd6",
   "metadata": {},
   "source": [
    "### Exercise 3.2\n",
    "Create a streamline plot for the case that the well discharge is $Q=300$ m$^3$/d. Start 20 streamlines at $x=-600$ m for $y$ varying from -300 to +300 m (as in the previous example). Next, add 20 streamlines that start at the river at $x=0$ and $y$ varying from -300 to +300 m. Some of the streamlines that start at the river will flow towards the well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d9bcf05-59ee-4060-835d-0677b819d577",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute grid of Qx, Qy\n",
    "Q = 300\n",
    "xg, yg = np.meshgrid(np.linspace(-600, 0, 100), np.linspace(-300, 300, 100))\n",
    "Qx, Qy = QxQy(xg, yg, xw, yw, Q, U)\n",
    "# plot streamlines\n",
    "plt.subplot(111, aspect=1)\n",
    "plt.axvline(0, color='skyblue', lw=8) # river\n",
    "plt.plot(xw, yw, 'k.')\n",
    "# Add your code below\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a14481-98ec-4b66-8641-2c7a15ec773b",
   "metadata": {},
   "source": [
    "### Capture zone when $Q>Q_\\text{max}$$^*$\n",
    "When the discharge of the well is larger than $Q_\\text{max}$, there are two stagnation points on the river. Between the stagnation points, water from the river flows towards the well. The locations of the stagnation points may be computed from:\n",
    "\\begin{equation}\n",
    "y_s = y_w \\pm \\sqrt{\\frac{-Qx_w}{\\pi U} - x_w^2}\n",
    "\\end{equation}\n",
    "A streamplot is shown for the case that $Q=300$ m$^3$/d, including the capture zone envelope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd97728-ac89-4cda-a505-7ab8fd70fe34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute grid of Qx, Qy\n",
    "Q = 300\n",
    "xg, yg = np.meshgrid(np.linspace(-600, 0, 100), np.linspace(-300, 300, 100))\n",
    "Qx, Qy = QxQy(xg, yg, xw, yw, Q, U)\n",
    "ys = yw + np.sqrt(-Q * xw / (np.pi * U) - xw ** 2)\n",
    "# plot streamlines\n",
    "plt.subplot(121, aspect=1)\n",
    "plt.axvline(0, color='skyblue', lw=8) # river\n",
    "xstart = -600 * np.ones(20)\n",
    "ystart = np.linspace(-300, 300, 20)\n",
    "for i in range(20):\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xstart[i], ystart[i])], \n",
    "                   integration_direction='forward')\n",
    "xstart = 0 * np.ones(20)\n",
    "ystart = np.linspace(-300, 300, 20)\n",
    "for i in range(20):\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(xstart[i], ystart[i])], \n",
    "                   integration_direction='forward', density=10)\n",
    "for y in [ys + 0.1, ys - 0.1, -ys + 0.1, -ys - 0.1]:\n",
    "    plt.streamplot(xg, yg, Qx, Qy, start_points=[(0, y)], color='k')\n",
    "plt.plot(xw, yw, 'k.');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "114fe7f6-10d3-47b5-a187-05fcd387e3a3",
   "metadata": {},
   "source": [
    "# Solution 4. Wells and rivers - the analytic element method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3473afde-47a1-4d58-8112-b2b76fd06dac",
   "metadata": {},
   "source": [
    "Pumping wells form an important part of groundwater models. The principle of superposition allows for the solution of problems with an arbitrary number of wells. The organization of solutions with multiple wells becomes cumbersome when the number of wells grows. The analytic element method is a modeling technique that facilitates the superposition of many solutions, including wells. \n",
    "\n",
    "The analytic element method is based on the superposition of analytic solutions. Each analytic solution represents a feature in the aquifer and has at least one free parameter. For example, the solution for a pumping well is an analytic element and the free parameter is the discharge of the well. The free parameter of an analytic element may either be specified or computed from specified conditions. For example, the discharge of a well may be computed from the head specified at the well (really, at the well radius). All unknown free parameters in a model are computed simultaneously as, e.g., the discharge of one well may effect the head at another well and vice versa. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47ce8fc6-7dca-4a8c-a70e-4a8cde61fef7",
   "metadata": {},
   "source": [
    "Here we use a relatively simple analytic element model that includes wells and stream segments. The analytic element code is included in the file `simple_aem.py` (a more sophisticated code written in Python is <A href=\"https://github.com/mbakker7/ttim\">TTIM</A>). As an example, we compute heads in an aquifer with three pumping wells as shown in the figure. In addition, the head is specified at one point called the reference point. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9193ca4-3c58-4c91-b206-645fd8634e95",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 200 # transmissivity, m^2/d\n",
    "x0, y0, Q0 = 0, 0, 400 # well 0\n",
    "x1, y1, Q1 = 200, 50, 600 # well 1\n",
    "x2, y2, Q2 = 100, 100, 200 # well 3\n",
    "rw = 0.3 # radius of the wells\n",
    "xr, yr, hr = 0, 200, 20 # head at reference point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c90ca36-0633-4327-9722-f21b82c6a0af",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "plt.plot([x0, x1, x2], [y0, y1, y2], 'k.')\n",
    "plt.text(x0, y0, f'Q={Q0} m$^3$/d', ha='center', va='bottom')\n",
    "plt.text(x1, y1, f'Q={Q1} m$^3$/d', ha='center', va='bottom')\n",
    "plt.text(x2, y2, f'Q={Q2} m$^3$/d', ha='center', va='bottom')\n",
    "plt.plot(xr, yr, 'x')\n",
    "plt.text(xr, yr, f'h={hr} m', ha='center', va='bottom')\n",
    "plt.grid()\n",
    "plt.xlim(-100, 300)\n",
    "plt.ylim(-100, 250);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e57d918-7f84-44bb-9fa4-bce6d1de1a4a",
   "metadata": {},
   "source": [
    "An analytic element is created as follows. First, the `Model`, `Well` and `Constant` classes are imported from the `simple_aem` class. Next, a model is created by specifying the transmissivity of the aquifer. The `Model` class returns a `Model` object, which is stored in the variable `ml` here (you can use any name you like). Next, an arbitratry number of `Well` elements may be added to the model. Finally, the head may be specified at one point with the `Constant` element. After all elements are added to the model, the model is solved and all unknown parameters are computed (in this example only the constant in the solution). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f438a5-0af0-4c5b-b324-b768437e8bcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/vincentpost/gwmodelling/main/simple_aem.py\n",
    "from simple_aem import Model, Well, Constant\n",
    "ml = Model(T)\n",
    "Well(ml, x0, y0, Q0, rw) # add well to model called ml\n",
    "Well(ml, x1, y1, Q1, rw)\n",
    "Well(ml, x2, y2, Q2, rw)\n",
    "Constant(ml, xr, yr, hr)\n",
    "ml.solve() # solve for unknowns in the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7139199d-7692-4861-ac3d-158b6fd6f233",
   "metadata": {},
   "source": [
    "The head may now be computed at any location with the `ml.head` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbbbb92a-7a86-45ff-89ce-43270f514a8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.head(100, 0) # head at (0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5352427c-8cf8-46df-95b1-9f057b0aa5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.head(0, 0) # head at well 0, automatically evaluated at well radius"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6d8e2c-e933-4429-b5a0-33cb084eee6d",
   "metadata": {},
   "source": [
    "A contour plot may be created, as before, by computing the head on a grid of points and using the `contour` routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3824305e-958f-46c6-ad25-65bdb0c296ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "xg, yg, = np.meshgrid(np.linspace(-100, 300, 100), np.linspace(-100, 200, 100))\n",
    "h = ml.head(xg, yg)\n",
    "print(f'min and max head: {h.min():.2f} m, {h.max():.2f} m')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113be17d-98c0-4709-9c53-bc8161167506",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(17, 21, 0.2))\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot([x0, x1, x2], [y0, y1, y2], 'k.');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc1cbf30-3eae-47ff-b7dd-66f1ba325de5",
   "metadata": {},
   "source": [
    "## Exercise 4.1\n",
    "Add a well to the model at $(x,y)=(100, 0)$. Determine the discharge of the well such that the head at the well is 18 m using trial and error. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b44779-ad18-4888-80e6-f0bfd6dd6a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d3ca3e-f6c5-41d2-8b54-948d10a394f1",
   "metadata": {},
   "source": [
    "## Wells with fixed head but unknown discharge\n",
    "In the previous example, the discharges of the wells were specified. What now if you are asked to pump the wells such that the head at each of the three wells is a specified value, for example 17 m? You can do that with the `HeadWell` element. For the `HeadWell` element you specify the location, the radius of the well, and the desired head at the well. The model then computes the discharge of the well (or the discharges of multiple wells) such that the heads are equal to the specified heads. In the example model below, the same three wells are used as before, but now the head is specified to be 17 m at all three wells, and the discharges of all wells are computed and stored. The discharge for each well is stored in the `.parameter` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feee70f4-157c-410b-8e21-9cc8e3904918",
   "metadata": {},
   "outputs": [],
   "source": [
    "from simple_aem import HeadWell\n",
    "ml = Model(T)\n",
    "w0 = HeadWell(ml, x0, y0, rw, 17) # add well to model called ml\n",
    "w1 = HeadWell(ml, x1, y1, rw, 17)\n",
    "w2 = HeadWell(ml, x2, y2, rw, 17)\n",
    "Constant(ml, xr, yr, hr)\n",
    "ml.solve() # solve for unknowns in the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbde1c29-df41-455f-be66-d4d3be4ff77c",
   "metadata": {},
   "outputs": [],
   "source": [
    "xg, yg, = np.meshgrid(np.linspace(-100, 300, 100), np.linspace(-100, 200, 100))\n",
    "h = ml.head(xg, yg)\n",
    "print(f'min and max head: {h.min()} m, {h.max()} m')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1d4ca2-b53a-4a3a-9b1f-8e3857e61eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(121, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(17, 21, 0.2))\n",
    "plt.clabel(cs, fmt='%1.1f')\n",
    "plt.plot([x0, x1, x2], [y0, y1, y2], 'k.');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c49562-1f4d-4c1f-89f8-6a9431d4f9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# computed discharge of well 0\n",
    "print(f'discharge of well 0: {w0.parameter:.0f} m^3/d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b794142-608d-47c7-8b4b-22ee32dd73d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print discharges of all wells:\n",
    "for i, w in enumerate([w0, w1, w2]):\n",
    "    print(f'dicharge of well {i} is {w.parameter:.0f} m^3/d')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bbe55df-ef2a-43cd-92ff-58428fde65f5",
   "metadata": {},
   "source": [
    "## Simulation of river segments with line-sinks$^*$\n",
    "River segments may be simulated with line-sinks. A line-sink is a line element that takes out water uniformly along a line segment. The head is specified at the center of the line-sink and the discharge of the line-sink is computed such that the head at the center of the segment is equal to the specfied value (very similar to a `HeadWell`). Line-sinks with a specified head are implemented in the `HeadLineSink` class, which is imported first. Next, a model is built for an area bounded by two river branches. Each river branch is divided in a number of segments. The water levels at the centers of the segments are indicated in the figure below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e297fa0-b916-41be-bb96-c4b7b6e36924",
   "metadata": {},
   "source": [
    "<img src=\"figs/aem_model.png\" width=\"400px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83c83f86-5a89-4c93-89a2-3d39efa99575",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "T = 100 # transmissivity, m^2/d\n",
    "# wells\n",
    "x0, y0, Q0 = 500, 250, 100 # well 0, discharge in m^3/d\n",
    "x1, y1, Q1 = 800, 500, 100 # well 1, discahrge in m^3/d\n",
    "Q1 = 100 # discharge, m^3/d\n",
    "# river segments, all in m\n",
    "xls0 = [0, 100, 200, 400, 600, 800, 1000, 1100, 1200]\n",
    "yls0 = [200, 200, 100, 100, 0, 0, 100, 300, 450]\n",
    "hls0 = np.linspace(39, 40.4, 8)\n",
    "xls1 = [0, 0, 200, 400, 600, 800, 1000, 1100, 1200]\n",
    "yls1 = [200, 400, 600, 600, 700, 700, 750, 800, 850]\n",
    "hls1 = np.linspace(39, 40.4, 8)\n",
    "# reference point\n",
    "xr, yr, hr = 0, 800, 39.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "976c4711-d09c-4dd6-b6ab-c5fe3724dbca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from simple_aem import HeadLineSink\n",
    "\n",
    "# analytic element model\n",
    "ml = Model(T) # transmissivity, m^2/d\n",
    "rf = Constant(ml, xr, yr, hr) # reference point\n",
    "w0 = Well(ml, x0, y0, Q0) # well 0\n",
    "w1 = Well(ml, x1, y1, Q1) # well 1\n",
    "for i in range(len(hls0)):\n",
    "    HeadLineSink(ml, xls0[i], yls0[i], xls0[i + 1], yls0[i + 1], hls0[i])\n",
    "for i in range(len(hls1)):\n",
    "    HeadLineSink(ml, xls1[i], yls1[i], xls1[i + 1], yls1[i + 1], hls1[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b21e18eb-952b-40d9-9663-918860696f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solution\n",
    "ml.solve()\n",
    "xg, yg = np.meshgrid(np.linspace(-100, 1300, 100), np.linspace(-100, 900, 100))\n",
    "h = ml.head(xg, yg)\n",
    "print(f'min and max head: {h.min():.2f} m, {h.max():.2f} m')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6883431-72cf-4b15-b71f-b31d436675f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# contour plot\n",
    "plt.subplot(111, aspect=1)\n",
    "cs = plt.contour(xg, yg, h, np.arange(38, 41, 0.1), colors='C0')\n",
    "plt.clabel(cs, fmt='%1.1f', fontsize='smaller')\n",
    "plt.plot(xls0, yls0, 'k')\n",
    "plt.plot(xls1, yls1, 'k')\n",
    "plt.plot([w0.zw.real, w1.zw.real], [w0.zw.imag, w1.zw.imag], 'k.');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b759c9b-7ffe-44a5-b7b5-2d79bc2ba1a5",
   "metadata": {},
   "source": [
    "Find more analytic solutions and their application in Python in [Bakker and Post, 2022, Analytical Groundwater Modeling](https://www.routledge.com/Analytical-Groundwater-Modeling-Theory-and-Applications-using-Python/Bakker-Post/p/book/9781138029392). "
   ]
  }
 ],
 "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.12"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}