{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cross-Sectional Models\n",
    "\n",
    "This notebook demonstrates how to build and solve cross-sectional (1D) models in `timflow`, including examples with strip inhomogeneities and infinitely long line elements."
   ]
  },
  {
   "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\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (4, 3)\n",
    "plt.rcParams[\"figure.autolayout\"] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two-layer model with head-specified line-sink\n",
    "Two-layer aquifer bounded on top by a semi-confined layer. Head above the semi-confining layer is 5. Head line-sink located at $x=0$ with head equal to 2, cutting through layer 0 only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelMaq(\n",
    "    kaq=[1, 2], z=[4, 3, 2, 1, 0], c=[1000, 1000], topboundary=\"semi\", hstar=5\n",
    ")\n",
    "ls = tfs.River1D(ml, xls=0, hls=2, layers=0)\n",
    "ml.solve()\n",
    "\n",
    "x = np.linspace(-200, 200, 101)\n",
    "h = ml.headalongline(x, np.zeros_like(x))\n",
    "plt.plot(x, h[0], label=\"layer 0\")\n",
    "plt.plot(x, h[1], label=\"layer 1\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1D inhomogeneity\n",
    "\n",
    "Three strips with semi-confined conditions on top of all three"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=2)\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-np.inf,\n",
    "    x2=-50,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=5,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-50,\n",
    "    x2=50,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=4.5,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=50,\n",
    "    x2=np.inf,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=4,\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
    "ml.plots.xsection(xy=[(-150, 0), (150, 0)], ax=ax, params=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-200, 200, 101)\n",
    "h = ml.headalongline(x, np.zeros(101))\n",
    "plt.plot(x, h[0], label=\"layer 0\")\n",
    "plt.plot(x, h[1], label=\"layer 1\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"head (m)\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.vcontoursf1D(x1=-200, x2=200, nx=100, levels=20, color=\"C0\", figsize=(10, 3));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Three strips with semi-confined conditions at the top of the strip in the middle only. The head is specified in the strip on the left and in the strip on the right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=2)\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-np.inf,\n",
    "    x2=-50,\n",
    "    kaq=[1, 2],\n",
    "    z=[3, 2, 1, 0],\n",
    "    c=[1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-50,\n",
    "    x2=50,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=4,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=50,\n",
    "    x2=np.inf,\n",
    "    kaq=[1, 2],\n",
    "    z=[3, 2, 1, 0],\n",
    "    c=[1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "rf1 = tfs.Constant(ml, -100, 0, 5)\n",
    "rf2 = tfs.Constant(ml, 100, 0, 5)\n",
    "\n",
    "ml.solve()\n",
    "\n",
    "ml.plots.xsection(xy=[(-100, 0), (100, 0)]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-200, 200, 101)\n",
    "h = ml.headalongline(x, np.zeros_like(x))\n",
    "Qx, _ = ml.disvecalongline(x, np.zeros_like(x))\n",
    "\n",
    "plt.figure(figsize=(10, 3))\n",
    "plt.subplot(121)\n",
    "plt.plot(x, h[0], label=\"layer 0\")\n",
    "plt.plot(x, h[1], label=\"layer 1\")\n",
    "plt.plot([-100, 100], [5, 5], \"k.\", label=\"fixed heads\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"head (m)\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.grid()\n",
    "plt.subplot(122)\n",
    "plt.plot(x, Qx[0], label=\"layer 0\")\n",
    "plt.plot(x, Qx[1], label=\"layer 1\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"$Q_x$ (m$^2$/d)\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.vcontoursf1D(x1=-200, x2=200, nx=100, levels=20, color=\"C0\", figsize=(10, 3));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Impermeable wall\n",
    "\n",
    "Flow from left to right in three-layer aquifer with impermeable wall in bottom 2 layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# need ModelMaq here since Uflow requires a confined background aquifer\n",
    "ml = tfs.ModelMaq(kaq=[1, 2, 4], z=[5, 4, 3, 2, 1, 0], c=[5000, 1000])\n",
    "uf = tfs.Uflow(ml, 0.002, 0)\n",
    "rf = tfs.Constant(ml, 100, 0, 20)\n",
    "ld1 = tfs.ImpermeableWall1D(ml, xld=0, layers=[0, 1])\n",
    "\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-100, 100, 101)\n",
    "h = ml.headalongline(x, np.zeros_like(x))\n",
    "Qx, _ = ml.disvecalongline(x, np.zeros_like(x))\n",
    "\n",
    "plt.figure(figsize=(10, 3))\n",
    "plt.subplot(121)\n",
    "plt.title(\"head\")\n",
    "plt.plot(x, h[0], label=\"layer 0\")\n",
    "plt.plot(x, h[1], label=\"layer 1\")\n",
    "plt.plot(x, h[2], label=\"layer 2\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"head (m)\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.grid()\n",
    "plt.subplot(122)\n",
    "plt.title(\"Qx\")\n",
    "plt.plot(x, Qx[0], label=\"layer 0\")\n",
    "plt.plot(x, Qx[1], label=\"layer 1\")\n",
    "plt.plot(x, Qx[2], label=\"layer 2\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"$Q_x$ (m$^2$/d)\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = ml.plots.vcontoursf1D(\n",
    "    x1=-200, x2=200, nx=100, levels=20, color=\"C0\", figsize=(10, 3), horizontal_axis=\"x\"\n",
    ")\n",
    "ld1.plot(ax);  # plot wall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infiltration\n",
    "\n",
    "Comparing solution with `Xsection` inhomogeneities to `XsectionAreaSink` solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=2)\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-np.inf,\n",
    "    x2=-50,\n",
    "    kaq=[1, 2],\n",
    "    z=[3, 2, 1, 0],\n",
    "    c=[1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-50,\n",
    "    x2=50,\n",
    "    kaq=[1, 2],\n",
    "    z=[3, 2, 1, 0],\n",
    "    c=[1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"conf\",\n",
    "    N=0.001,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=50,\n",
    "    x2=np.inf,\n",
    "    kaq=[1, 2],\n",
    "    z=[3, 2, 1, 0],\n",
    "    c=[1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "tfs.Constant(ml, -100, 0, 10)\n",
    "tfs.Constant(ml, 100, 0, 10)\n",
    "ml.solve()\n",
    "\n",
    "ml.plots.vcontoursf1D(x1=-100, x2=100, nx=100, levels=20, color=\"C0\", figsize=(10, 3));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml2 = tfs.ModelMaq(kaq=[1, 2], z=[3, 2, 1, 0], c=[1000], topboundary=\"conf\")\n",
    "tfs.XsectionAreaSink(ml2, -50, 50, 0.001)\n",
    "tfs.Constant(ml2, -100, 0, 10)\n",
    "ml2.solve()\n",
    "ml2.plots.vcontoursf1D(\n",
    "    x1=-100, x2=100, nx=100, levels=20, color=\"C0\", figsize=(10, 3), horizontal_axis=\"x\"\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-100, 100, 100)\n",
    "plt.plot(x, ml.headalongline(x, 0)[0], \"C0\")\n",
    "plt.plot(x, ml.headalongline(x, 0)[1], \"C0\")\n",
    "plt.plot(x, ml2.headalongline(x, 0)[0], \"--C1\")\n",
    "plt.plot(x, ml2.headalongline(x, 0)[1], \"--C1\")\n",
    "plt.xlabel(\"x (m)\")\n",
    "plt.ylabel(\"head (m)\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=50)\n",
    "tfs.Xsection3D(ml, x1=-np.inf, x2=-5, kaq=1, z=np.arange(5, -0.1, -0.1), kzoverkh=0.1)\n",
    "tfs.Xsection3D(\n",
    "    ml,\n",
    "    x1=-5,\n",
    "    x2=5,\n",
    "    kaq=1,\n",
    "    z=np.arange(5, -0.1, -0.1),\n",
    "    kzoverkh=0.1,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=5.5,\n",
    "    topres=3,\n",
    "    topthick=0.3,\n",
    ")\n",
    "tfs.Xsection3D(ml, x1=5, x2=np.inf, kaq=1, z=np.arange(5, -0.1, -0.1), kzoverkh=0.1)\n",
    "rf1 = tfs.Constant(ml, -100, 0, 5.7)\n",
    "rf2 = tfs.Constant(ml, 100, 0, 5.47)\n",
    "\n",
    "ml.solve()\n",
    "\n",
    "ml.plots.vcontoursf1D(x1=-20, x2=20, nx=100, levels=20, color=\"C0\", figsize=(10, 3));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=5)\n",
    "tfs.Xsection3D(ml, x1=-np.inf, x2=-5, kaq=1, z=np.arange(5, -0.1, -1), kzoverkh=0.1)\n",
    "tfs.Xsection3D(\n",
    "    ml,\n",
    "    x1=-5,\n",
    "    x2=5,\n",
    "    kaq=1,\n",
    "    z=np.arange(5, -0.1, -1),\n",
    "    kzoverkh=0.1,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=5.5,\n",
    "    topres=3,\n",
    "    topthick=0.3,\n",
    ")\n",
    "tfs.Xsection3D(ml, x1=5, x2=np.inf, kaq=1, z=np.arange(5, -0.1, -1), kzoverkh=0.1)\n",
    "rf1 = tfs.Constant(ml, -100, 0, 5.7)\n",
    "rf2 = tfs.Constant(ml, 100, 0, 5.47)\n",
    "\n",
    "ml.solve()\n",
    "\n",
    "ml.plots.vcontoursf1D(x1=-20, x2=20, nx=100, levels=20, color=\"C0\", figsize=(10, 3));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelXsection(naq=2)\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-np.inf,\n",
    "    x2=-50,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=15,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=-50,\n",
    "    x2=50,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=13,\n",
    ")\n",
    "tfs.XsectionMaq(\n",
    "    ml,\n",
    "    x1=50,\n",
    "    x2=np.inf,\n",
    "    kaq=[1, 2],\n",
    "    z=[4, 3, 2, 1, 0],\n",
    "    c=[1000, 1000],\n",
    "    npor=0.3,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=11,\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from timflow.steady.linesink1d import FluxDiffLineSink1D, HeadDiffLineSink1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tfs.ModelMaq(kaq=[10], z=[0, -10], topboundary=\"conf\")\n",
    "ls = tfs.River1D(ml, xls=0, hls=1, wh=\"H\", layers=0)\n",
    "ls = tfs.River1D(ml, xls=200, hls=0, wh=\"H\", layers=0)\n",
    "hd = HeadDiffLineSink1D(ml, xls=100)\n",
    "fd = FluxDiffLineSink1D(ml, xls=100)\n",
    "ml.solve()\n",
    "\n",
    "x = np.linspace(0, 200, 101)\n",
    "h = ml.headalongline(x, np.zeros_like(x))\n",
    "plt.plot(x, h[0], label=\"layer 0\")\n",
    "# plt.plot(x, h[1], label=\"layer 1\")\n",
    "plt.legend(loc=\"best\");"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}