{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connected Wells\n",
    "\n",
    "WellString elements can be used to model a series of connected wells. There are several types of\n",
    "WellString:\n",
    "\n",
    "- `WellString`: connected wells with equal head inside the wells and a user-specified total discharge\n",
    "- `HeadWellString`: connected wells with equal specified head inside the wells, this is identical to specifying multiple HeadWells.\n",
    "- `TargetHeadWellString`: connected wells with equal head inside the wells and one user-specified head at a point $(x_c, y_c)$\n",
    "\n",
    "This notebook shows how to model the `WellString` and `TargetHeadWellString` elements\n",
    "and compares the results to models with individual wells. Examples are shown for both\n",
    "elements in single- and multi-layer models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import timflow.steady as tfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WellString"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A series of wells that have equal head (pressure) inside the well and pump with a total specified discharge."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single-layer model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model parameters\n",
    "kh = 10  # m/day\n",
    "\n",
    "ctop = 1000.0  # resistance top leaky layer in days\n",
    "\n",
    "ztop = 0.0  # surface elevation\n",
    "zbot = -20.0  # bottom elevation of the model\n",
    "\n",
    "z = [1.0, ztop, zbot]\n",
    "kaq = np.array([kh])\n",
    "c = np.array([ctop])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reference model with 2 wells with equal discharge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlref0 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "w1 = tfs.Well(mlref0, 0, -10, Qw=50, rw=0.1)\n",
    "w2 = tfs.Well(mlref0, 0, 10, Qw=50, rw=0.1)\n",
    "mlref0.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reference model with 2 head-specified wells, specifying the head inside well `w1` computed with the previous model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlref1 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "wh1 = tfs.HeadWell(mlref1, 0, -10, hw=w1.headinside().item(), rw=0.1)\n",
    "wh2 = tfs.HeadWell(mlref1, 0, 10, hw=w1.headinside().item(), rw=0.1)\n",
    "mlref1.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model with a WellString that has total discharge equal to the sum of the two discharge-specified wells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml0 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.WellString(ml0, xy=[(0, -10), (0, 10)], Qw=100, rw=0.1)\n",
    "ml0.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Contour the heads of the first reference model and the WellString model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels = 10\n",
    "ax = mlref0.plots.topview(win=(-50, 50, -50, 50))\n",
    "mlref0.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[0],\n",
    "    ax=ax,\n",
    ")\n",
    "ml0.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[0],\n",
    "    linestyles=\"dashed\",\n",
    "    color=\"C1\",\n",
    "    ax=ax,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare heads along y between all 3 models (2 reference models and the WellString model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = np.zeros_like(y)\n",
    "href = mlref0.headalongline(x, y)\n",
    "href2 = mlref1.headalongline(x, y)\n",
    "h0 = ml0.headalongline(x, y)\n",
    "\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(y, href[0], label=\"Reference Wells\")\n",
    "plt.plot(y, href2[0], \"--\", label=\"Reference HeadWells\")\n",
    "plt.plot(y, h0[0], \"k.\", ms=3, label=\"WellString model\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.ylabel(\"head [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check computed discharges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 discharge wells, Q specified not computed\n",
    "w1.discharge(), w2.discharge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 HeadWells\n",
    "wh1.discharge(), wh2.discharge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WellString\n",
    "ws.discharge_per_well()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare heads inside the wells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 discharge wells\n",
    "w1.headinside(), w2.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 HeadWells, specified, not computed\n",
    "wh1.headinside(), wh2.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WellString\n",
    "ws.headinside(), ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multilayer model\n",
    "\n",
    "The following example compares the WellString element to reference models with\n",
    "individual Wells and HeadWells in an aquifersystem consisting of 3 layers. In this example, all wells are screened in the bottom two aquifers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model parameters\n",
    "kh = [5, 10, 20]  # m/day\n",
    "\n",
    "c = [1000.0, 100.0, 1.0]  # resistance leaky layers in days\n",
    "\n",
    "ztop = 0.0  # surface elevation\n",
    "zbot = -50.0  # bottom elevation of the model\n",
    "\n",
    "z = [1.0, ztop, -10, -15, -25, -26, zbot]\n",
    "kaq = np.array(kh)\n",
    "c = np.array(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reference model with discharge wells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlref2 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "w1 = tfs.Well(mlref2, 0, -10, Qw=50, rw=0.1, layers=[1, 2])\n",
    "w2 = tfs.Well(mlref2, 0, 10, Qw=50, rw=0.1, layers=[1, 2])\n",
    "mlref2.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reference model with head-specified wells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlref3 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "wh1 = tfs.HeadWell(mlref3, 0, -10, hw=w1.headinside(), rw=0.1, layers=[1, 2])\n",
    "wh2 = tfs.HeadWell(mlref3, 0, 10, hw=w1.headinside(), rw=0.1, layers=[1, 2])\n",
    "mlref3.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model with WellString"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml1 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.WellString(ml1, xy=[(0, -10), (0, 10)], Qw=100, rw=0.1, layers=[1, 2])\n",
    "ml1.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare head contours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ilay = 1\n",
    "levels = 10\n",
    "ax = mlref2.plots.topview(win=(-50, 50, -50, 50))\n",
    "mlref2.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[ilay],\n",
    "    ax=ax,\n",
    ")\n",
    "ml1.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[ilay],\n",
    "    linestyles=\"dashed\",\n",
    "    color=\"C1\",\n",
    "    ax=ax,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare drawdowns along y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = np.zeros_like(y)\n",
    "href2 = mlref2.headalongline(x, y)\n",
    "href3 = mlref3.headalongline(x, y)\n",
    "h1 = ml1.headalongline(x, y)\n",
    "\n",
    "ilay = 1\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(y, href2[ilay], label=\"Reference Well\")\n",
    "plt.plot(y, href3[ilay], \"--\", label=\"Reference HeadWell\")\n",
    "plt.plot(y, h1[ilay], \"k.\", ms=3, label=\"WellString model\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.ylabel(\"head [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare discharges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 discharge wells, total Q specified\n",
    "w1.discharge(), w2.discharge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 head specified wells\n",
    "wh1.discharge(), wh2.discharge()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WellString, transposed to compare to above\n",
    "ws.discharge_per_well().T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 discharge wells\n",
    "w1.headinside(), w2.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 HeadWells, specified, not computed\n",
    "wh1.headinside(), wh2.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.headinside(), ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TargetHeadWellString\n",
    "\n",
    "A series of connected wells with equal head (pressure) inside all wells and a specified\n",
    "head at one point $(x_c, y_c)$ and layer. By default the specified head is applied in layer\n",
    "0, but this can be set with `lcp=<layer>`in the `TargetHeadWellString` element."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single layer example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model parameters\n",
    "kh = 10  # m/day\n",
    "\n",
    "ctop = 1000.0  # resistance top leaky layer in days\n",
    "\n",
    "ztop = 0.0  # surface elevation\n",
    "zbot = -20.0  # bottom elevation of the model\n",
    "\n",
    "z = [1.0, ztop, zbot]\n",
    "kaq = np.array([kh])\n",
    "c = np.array([ctop])\n",
    "\n",
    "# point at which head is specified\n",
    "xcp = 10.0\n",
    "ycp = 0.0\n",
    "hcp = -2.0  # specified head at (xc, yc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a model with a `TargetHeadWellString`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml2 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.TargetHeadWellString(ml2, [(0, -10), (0, 10)], rw=0.1, hcp=hcp, xcp=10, ycp=0)\n",
    "ml2.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the head contours for the model and mark location of the head target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ilay = 0\n",
    "levels = [-2.4, -2.2, -2.0, -1.8, -1.6, -1.4, -1.2]\n",
    "ml2.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[ilay],\n",
    "    newfig=False,\n",
    "    linestyles=\"dashed\",\n",
    "    color=\"C1\",\n",
    ")\n",
    "plt.plot(ws.xcp, ws.ycp, \"kx\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the head along y at x=10, to show that the specified head is met"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = ws.xcp * np.ones_like(y)\n",
    "h2 = ml2.headalongline(x, y)\n",
    "\n",
    "ilay = 0\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(y, h2[ilay], label=\"TargetHeadWellString model\")\n",
    "plt.plot(ws.ycp, ws.hcp, \"kx\", label=\"Specified head\")\n",
    "plt.ylabel(\"head (at x=10) [m]\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the discharge per well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TargetHeadWellString model, shape: (nlay, nwells)\n",
    "ws.discharge_per_well()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare heads inside the well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TargetHeadWellString\n",
    "ws.headinside(), ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-layer example\n",
    "\n",
    "An example of a TargetHeadWellString in a multi-layer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model parameters\n",
    "kh = [5, 10, 20]  # m/day\n",
    "\n",
    "c = [1000.0, 100.0, 1.0]  # resistance leaky layers in days\n",
    "\n",
    "ztop = 0.0  # surface elevation\n",
    "zbot = -50.0  # bottom elevation of the model\n",
    "\n",
    "z = [1.0, ztop, -10, -15, -25, -26, zbot]\n",
    "kaq = np.array(kh)\n",
    "c = np.array(c)\n",
    "\n",
    "# point at which head is specified\n",
    "xcp = 10.0\n",
    "ycp = 0.0\n",
    "hcp = -2.0  # specified head at (xc, yc)\n",
    "lcp = 1  # layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a model with a TargetHeadWellString."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml3 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.TargetHeadWellString(\n",
    "    ml3,\n",
    "    xy=[(0, -10), (0, 10)],\n",
    "    rw=0.1,\n",
    "    layers=[0, 1, 2],\n",
    "    hcp=hcp,\n",
    "    xcp=xcp,\n",
    "    ycp=ycp,\n",
    "    lcp=lcp,\n",
    ")\n",
    "ml3.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the head contours in the reference model and the TargetHeadWellString model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ilay = 1\n",
    "levels = [-2.4, -2.2, -2.0, -1.8, -1.6, -1.4, -1.2]\n",
    "ml3.plots.contour(\n",
    "    win=(-50, 50, -50, 50),\n",
    "    ngr=51,\n",
    "    levels=levels,\n",
    "    decimals=2,\n",
    "    layers=[ilay],\n",
    "    newfig=False,\n",
    "    linestyles=\"dashed\",\n",
    "    color=\"C1\",\n",
    ")\n",
    "plt.plot(ws.xcp, ws.ycp, \"kx\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the head along y=10 to show that the head target is met."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = ws.xcp * np.ones_like(y)\n",
    "h3 = ml3.headalongline(x, y)\n",
    "\n",
    "ilay = lcp\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(y, h3[ilay], label=\"TargetHeadWellString model\")\n",
    "plt.plot(ws.ycp, ws.hcp, \"kx\", label=\"Specified head\")\n",
    "plt.ylabel(\"head (at x=10) [m]\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare discharges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TargetHeadWellString, discharge, shape : (nlay, nwells)\n",
    "ws.discharge_per_well()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare heads inside the well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TargetHeadWellString\n",
    "ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute head at specified point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TargetHeadWellString\n",
    "ml3.head(xcp, ycp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Different layers per well\n",
    "\n",
    "The WellString elements allow the specification of different layers per well, as shown\n",
    "in the example below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml4 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.TargetHeadWellString(\n",
    "    ml4, xy=[(0, -10), (0, 10)], rw=0.1, layers=[(1,), (2,)], hcp=-1, xcp=10, ycp=0\n",
    ")\n",
    "ml4.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the wells in a cross-section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 3))\n",
    "ml4.plots.xsection(xy=[(0, 20), (0, -20)], labels=True, ax=ax)\n",
    "for w in ws.wlist:\n",
    "    for ilay in w.layers:\n",
    "        ax.plot([w.yc, w.yc], [ml4.aq.zaqtop[ilay], ml4.aq.zaqbot[ilay]], \"k-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a cross section of the head along y, showing all layers, and showing that the head\n",
    "in layer 0 runs through -1, the head we specified in the TargetHeadWellString."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = ws.xcp * np.ones_like(y)\n",
    "h4 = ml4.headalongline(x, y)\n",
    "\n",
    "ilay = 0\n",
    "plt.figure(figsize=(10, 2))\n",
    "for ilay in range(ml4.aq.naq):\n",
    "    plt.plot(y, h4[ilay], \"-\", label=f\"layer {ilay}\")\n",
    "plt.plot(ws.ycp, ws.hcp, \"kx\", label=\"Specified head in layer 0\")\n",
    "plt.ylabel(\"head (at x=10) [m]\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=4);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get total well discharge per layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.discharge()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or per well, per layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.discharge_per_well()  # shape : (nlay, nwells)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show that head inside the wells is equal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `TargetHeadWellString` element also has a `headinside()` method that returns a\n",
    "single value (since the head inside must be equal)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test first well screened in two layers\n",
    "\n",
    "Test a TargetHeadWellString when the first well is screened in two layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lcp = 1\n",
    "ml5 = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "ws = tfs.TargetHeadWellString(\n",
    "    ml5,\n",
    "    xy=[(0, -10), (0, 10)],\n",
    "    rw=0.1,\n",
    "    layers=[(0, 1), (2,)],\n",
    "    res=0.0,\n",
    "    hcp=-1,\n",
    "    xcp=10,\n",
    "    ycp=0,\n",
    "    lcp=lcp,\n",
    ")\n",
    "ml5.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 3))\n",
    "ml5.plots.xsection(xy=[(0, 20), (0, -20)], labels=True, ax=ax)\n",
    "for w in ws.wlist:\n",
    "    for ilay in w.layers:\n",
    "        ax.plot([w.yc, w.yc], [ml5.aq.zaqtop[ilay], ml5.aq.zaqbot[ilay]], \"k-\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.linspace(-50, 50, 101)\n",
    "x = ws.xcp * np.ones_like(y)\n",
    "h5 = ml5.headalongline(x, y)\n",
    "\n",
    "ilay = lcp\n",
    "plt.figure(figsize=(10, 2))\n",
    "for ilay in range(ml5.aq.naq):\n",
    "    plt.plot(y, h5[ilay], \"-\", label=f\"layer {ilay}\")\n",
    "plt.plot(ws.ycp, ws.hcp, \"kx\", label=\"Specified head in layer 0\")\n",
    "plt.ylabel(\"head (at x=10) [m]\")\n",
    "plt.xlabel(\"y [m]\")\n",
    "plt.legend(loc=(0, 1), frameon=False, ncol=4);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the heads inside the wells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.wlist[0].headinside(), ws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ws.headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the head at the target point, this should be equal to the specified head in the layer we entered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml5.head(xcp, ycp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HeadWellString with resistance\n",
    "\n",
    "Test the result when we add resistance to the wells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "hws = tfs.HeadWellString(ml, xy=[(0, -10), (0, 10)], rw=0.1, res=0.1, layers=[0], hw=-2)\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The head inside the well should be equal to the specified head."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hws.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hws.wlist[0].headinside(), hws.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TargetHeadWellString with resistance\n",
    "\n",
    "Test a `TargetHeadWellString` with resistance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lcp = 1\n",
    "ml = tfs.ModelMaq(kaq=kaq, c=c, z=z, topboundary=\"semi\", hstar=0)\n",
    "thw = tfs.TargetHeadWellString(\n",
    "    ml,\n",
    "    xy=[(0, -10), (0, 10)],\n",
    "    rw=0.1,\n",
    "    layers=[(1,), (2,)],\n",
    "    res=0.01,\n",
    "    hcp=-1,\n",
    "    xcp=10,\n",
    "    ycp=0,\n",
    "    lcp=lcp,\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that the head inside the well is equal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thw.headinside()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thw.wlist[0].headinside(), thw.wlist[1].headinside()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that the head at the target point is met."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.head(thw.xcp, thw.ycp)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}