{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test integrated normal flux\n",
    "\n",
    "This notebook demonstrates and tests the `intnormalflux` method for `timflow` models.\n",
    "\n",
    "This method integrates the flux normal to a line along that line. Two integration methods are implemented:\n",
    "- numerical integration using `scipy.optimize.quad_vec`\n",
    "- analytic approximation of the integral using Legendre polynomials\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import timflow.steady as tfs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uniform flow field\n",
    "\n",
    "First, we perform some simple sanity checks in a uniform flow field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slope = 0.001  # head gradient in m/m\n",
    "\n",
    "ml = tfs.ModelMaq()\n",
    "u = tfs.Uflow(ml, slope, angle=0.0)\n",
    "c = tfs.Constant(ml, 0.0, 0.0, 10)\n",
    "ml.solve()\n",
    "ml.plots.contour([-1000, 1000, -1000, 1000], decimals=1);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a line between $(-10, -10)$ and $(10, 10)$ which has a 45 degree angle relative to the positive x-axis. The integrated flux along this line should equal $L \\cdot \\cos(\\theta) \\cdot \\text{slope}$.\n",
    "\n",
    "Since we have defined Q as positive when flowing to the left when going from \n",
    "$(x_1, y_1)$ to $(x_2, y_2)$, `intfluxnorm` will return a negative flux, but the absolute value should match our calculation. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, y1 = -10.0, -10.0\n",
    "x2, y2 = 10.0, 10.0\n",
    "L = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)\n",
    "theta = np.arctan2(y2 - y1, x2 - x1)\n",
    "\n",
    "qn_quad = ml.intnormflux_segment(x1, y1, x2, y2, method=\"quad\")\n",
    "qn_leg = ml.intnormflux_segment(x1, y1, x2, y2, method=\"legendre\", ndeg=3)\n",
    "\n",
    "print(np.round(np.array([qn_quad[0], qn_leg[0], L * np.cos(theta) * slope]), 4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we define a line normal to the head gradient. The integrated normal flux along that should equal the length of the line multiplied by the gradient: $L \\cdot \\text{slope}$. Once again our integrated flux is negative since we define flow to the left as positive when integrating from $(x_1, y_1)$ to $(x_2, y_2)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, y1 = 0.0, -10.0\n",
    "x2, y2 = 0.0, 10.0\n",
    "L = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)\n",
    "\n",
    "qn_quad = ml.intnormflux_segment(x1, y1, x2, y2, method=\"quad\")\n",
    "qn_leg = ml.intnormflux_segment(x1, y1, x2, y2, method=\"legendre\", ndeg=3)\n",
    "\n",
    "print(np.round(np.array([qn_quad[0], qn_leg[0], L * slope]), 4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a line parallel to the head gradient the integrated flux normal to the line should equal 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, y1 = -10.0, 0.0\n",
    "x2, y2 = 10.0, 0.0\n",
    "\n",
    "qn_quad = ml.intnormflux_segment(x1, y1, x2, y2, method=\"quad\")\n",
    "qn_leg = ml.intnormflux_segment(x1, y1, x2, y2, method=\"legendre\", ndeg=3)\n",
    "\n",
    "print(np.round(np.array([qn_quad[0], qn_leg[0], 0.0]), 4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two layer confined model with well\n",
    "\n",
    "In this next example we create a model with two layers with a confined top and a well in the bottom layer that pumps 500 $m^3/d$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelMaq(kaq=[5.0, 10.0], z=[0, -10, -15, -30], c=[100])\n",
    "c = tfs.Constant(ml, 1e4, 1e4, 10)\n",
    "w = tfs.Well(ml, 0.0, 0.0, Qw=500, layers=[1])\n",
    "ml.solve()\n",
    "ml.plots.contour(\n",
    "    [-1000, 1000, -1000, 1000],\n",
    "    levels=np.arange(6.0, 9.5, 0.1),\n",
    "    ngr=51,\n",
    "    layers=[0, 1],\n",
    "    decimals=1,\n",
    ");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the head along radial distance $r$ for both layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xl = np.linspace(0.1, 1000, 301)\n",
    "yl = np.zeros_like(xl)\n",
    "\n",
    "h = ml.headalongline(xl, yl)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
    "ax.plot(xl, h[0], color=\"C0\")\n",
    "ax.plot(xl, h[1], color=\"C1\")\n",
    "ax.set_ylabel(\"head (m)\")\n",
    "ax.grid(True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define two arbitrary polygons along which we will calculate the total net inflow using `intnormflux`. One polygon contains the well, the other does not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xy_in = [\n",
    "    (-200, -100),\n",
    "    (-25, -250),\n",
    "    (200, -125),\n",
    "    (150, 75),\n",
    "    (0, 120),\n",
    "    (-125, 50),\n",
    "    (-200, -100),\n",
    "]\n",
    "\n",
    "xy_out = [\n",
    "    (250, 0),\n",
    "    (290, 90),\n",
    "    (280, 110),\n",
    "    (210, 100),\n",
    "    (200, 30),\n",
    "    (250, 0),\n",
    "]\n",
    "\n",
    "window = [-250, 300, -300, 175]\n",
    "ml.plots.topview(window)\n",
    "for i in range(len(xy_in) - 1):\n",
    "    xyi = np.array(xy_in[i : i + 2])\n",
    "    (p1,) = plt.plot(xyi[:, 0], xyi[:, 1], ls=\"dashed\", color=\"k\", label=\"well inside\")\n",
    "for i in range(len(xy_out) - 1):\n",
    "    xyi = np.array(xy_out[i : i + 2])\n",
    "    (p2,) = plt.plot(xyi[:, 0], xyi[:, 1], ls=\"dashed\", color=\"r\", label=\"well outside\")\n",
    "plt.legend([p1, p2], [p1.get_label(), p2.get_label()], loc=(0, 1), frameon=False, ncol=2)\n",
    "plt.grid(True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Integration of the normal flux along edges of the polygon with the well inside shows that the total inflow along the shape is equal to $Q_w$. Both integration methods (quad and legendre) yield similar results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Qn_quad = ml.intnormflux(xy_in, method=\"quad\")\n",
    "Qn_leg = ml.intnormflux(xy_in, method=\"legendre\", ndeg=7)\n",
    "Qn_def = ml.intnormflux(xy_in)\n",
    "\n",
    "print(\"Q (quad)    :\", Qn_quad.sum())\n",
    "print(\"Q (legendre):\", Qn_leg.sum())\n",
    "print(\"Q (default):\", Qn_def.sum())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Integration of the normal flux along the edges of the polygon with no well inside equals 0. Any water flowing in must by definition also flow out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Qn_quad = ml.intnormflux(xy_out, method=\"quad\")\n",
    "Qn_leg = ml.intnormflux(xy_out, method=\"legendre\", ndeg=7)\n",
    "Qn_def = ml.intnormflux(xy_out)\n",
    "\n",
    "print(\"Q (quad)    :\", Qn_quad.sum())\n",
    "print(\"Q (legendre):\", Qn_leg.sum())\n",
    "print(\"Q (default):\", Qn_def.sum())"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}