{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HeadDiffLineSink1D and FluxDiffLineSink1D\n",
    "\n",
    "Continuity of head and continuity of flow."
   ]
  },
  {
   "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",
    "import timflow.transient as tft\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 10.0\n",
    "H = 10.0\n",
    "S = 0.2\n",
    "\n",
    "delh = 2.0\n",
    "t0 = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlconf = tft.ModelMaq(\n",
    "    kaq=k, z=[0, -H], Saq=S, tmin=1, tmax=1e2, topboundary=\"conf\", phreatictop=True\n",
    ")\n",
    "hls_left = tft.River1D(mlconf, xls=0.0, tsandh=[(0, delh)], layers=[0])\n",
    "hls_right = tft.River1D(mlconf, xls=200.0, tsandh=[(0, 0.0)], layers=[0])\n",
    "\n",
    "# headdiff on right side, fluxdiff on left side\n",
    "hdiff = tft.HeadDiffLineSink1D(mlconf, xls=100.0 + 1e-12, layers=[0])\n",
    "qdiff = tft.FluxDiffLineSink1D(mlconf, xls=100.0 - 1e-12, layers=[0])\n",
    "\n",
    "mlconf.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 200, 101)\n",
    "y = np.zeros_like(x)\n",
    "t = np.logspace(0, 2, 3)\n",
    "\n",
    "for i in range(len(t)):\n",
    "    h = mlconf.headalongline(x, y, t[i])\n",
    "    plt.plot(x, h.squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(\"x [m]\")\n",
    "plt.ylabel(\"head [m]\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-aquifer cross-sections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = [0.0, -5.0, -6.0, -15.0]\n",
    "Saq = [0.2, 1e-4]\n",
    "c = [10.0]\n",
    "k = np.array([5.0, 10.0])\n",
    "k_left = 0.5 * k\n",
    "k_right = k\n",
    "delh = 1.0\n",
    "res = 10.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transient model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlconf2 = tft.ModelXsection(naq=2, tmin=1, tmax=1e2)\n",
    "inhom_left = tft.XsectionMaq(\n",
    "    mlconf2,\n",
    "    -np.inf,\n",
    "    100.0,\n",
    "    kaq=k_left,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "inhom_right = tft.XsectionMaq(\n",
    "    mlconf2,\n",
    "    100.0,\n",
    "    np.inf,\n",
    "    kaq=k_right,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "\n",
    "hls_left = tft.River1D(mlconf2, xls=0.0, tsandh=[(0, delh)], layers=[0, 1])\n",
    "hls_right = tft.River1D(mlconf2, xls=200.0, tsandh=[(0, 0.0)], layers=[0, 1])\n",
    "\n",
    "mlconf2.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Steady-state model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlss = tfs.ModelMaq(kaq=k, z=z, c=c, topboundary=\"conf\")\n",
    "inhom_left = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    -np.inf,\n",
    "    100.0,\n",
    "    kaq=k_left,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "inhom_right = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    100.0,\n",
    "    np.inf,\n",
    "    kaq=k_right,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "\n",
    "hls_left = tfs.River1D(mlss, xls=0.0, hls=1.0, layers=[0, 1])\n",
    "hls_right = tfs.River1D(mlss, xls=200.0, hls=0.0, layers=[0, 1])\n",
    "\n",
    "mlss.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 200, 101)\n",
    "y = np.zeros_like(x)\n",
    "t = np.logspace(0, 2, 3)\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(10, 5))\n",
    "\n",
    "for i in range(len(t)):\n",
    "    h = mlconf2.headalongline(x, y, t[i])\n",
    "    ax0.plot(x, h[0].squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "    ax1.plot(x, h[1].squeeze(), label=f\"t={t[i]:.0f} d\", ls=\"dashed\")\n",
    "\n",
    "hss = mlss.headalongline(x, y)\n",
    "ax0.plot(x, hss[0].squeeze(), lw=1.5, color=\"k\", ls=\"dotted\", label=\"timflow\")\n",
    "ax1.plot(x, hss[1].squeeze(), lw=1.5, color=\"k\", ls=\"dotted\")\n",
    "\n",
    "ax0.legend(loc=(0, 1), frameon=False, ncol=2)\n",
    "ax0.set_title(\"Layer 0\")\n",
    "ax1.set_xlabel(\"x [m]\")\n",
    "ax1.set_title(\"Layer 1\")\n",
    "for ax in [ax0, ax1]:\n",
    "    ax.set_ylabel(\"head [m]\")\n",
    "    ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = 1e-6\n",
    "t = np.logspace(0, 2, 100)\n",
    "hleft = mlconf2.head(100 - dx, 0, t)\n",
    "hright = mlconf2.head(100 + dx, 0, t)\n",
    "disxleft, _ = mlconf2.disvec(100 - dx, 0, t)\n",
    "disxright, _ = mlconf2.disvec(100 + dx, 0, t)\n",
    "disxnum = inhom_left.Haq[:, np.newaxis] * (hleft - hright) / res\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "for i in range(2):\n",
    "    plt.subplot(2, 1, i + 1)\n",
    "    plt.semilogx(t, hleft[i], label=\"hleft\")\n",
    "    plt.semilogx(t, hright[i], \"--\", label=\"hright\")\n",
    "    plt.semilogx(t, disxleft[i], label=\"Qxleft timflow\")\n",
    "    plt.semilogx(t, disxright[i], \"--\", label=\"Qxright timflow\")\n",
    "    if i == 0:  # otherwise no leaky wall\n",
    "        plt.semilogx(t, disxnum[i], \"k:\", label=\"Qx numerical\")\n",
    "    plt.xlabel(\"time (d)\")\n",
    "    plt.ylabel(\"head or Qx\")\n",
    "    plt.title(f\"layer {i}\")\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 strip inhomogeneities\n",
    "\n",
    "Test if 3 strip inhomogeneities are working. Introduce a middle inhom with hydraulic conductivity $2k$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_mid = [2 * k]\n",
    "\n",
    "mlconf2 = tft.ModelXsection(naq=2, tmin=1, tmax=1e2)\n",
    "\n",
    "inhom_left = tft.XsectionMaq(\n",
    "    mlconf2,\n",
    "    -np.inf,\n",
    "    50.0,\n",
    "    kaq=k_left,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "inhom_mid = tft.XsectionMaq(\n",
    "    mlconf2,\n",
    "    50,\n",
    "    150.0,\n",
    "    kaq=2 * k,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "inhom_right = tft.XsectionMaq(\n",
    "    mlconf2,\n",
    "    150.0,\n",
    "    np.inf,\n",
    "    kaq=k_right,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "\n",
    "hls_left = tft.River1D(mlconf2, xls=0.0, tsandh=[(0, delh)], layers=[0, 1])\n",
    "hls_right = tft.River1D(mlconf2, xls=200.0, tsandh=[(0, 0.0)], layers=[0, 1])\n",
    "\n",
    "mlconf2.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 200, 101)\n",
    "y = np.zeros_like(x)\n",
    "t = np.logspace(0, 2, 3)\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(10, 5))\n",
    "\n",
    "for i in range(len(t)):\n",
    "    h = mlconf2.headalongline(x, y, t[i])\n",
    "    ax0.plot(x, h[0].squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "    ax1.plot(x, h[1].squeeze(), label=f\"t={t[i]:.0f} d\", ls=\"dashed\")\n",
    "\n",
    "ax0.legend(loc=(0, 1), frameon=False, ncol=3)\n",
    "ax0.set_title(\"Layer 0\")\n",
    "ax1.set_xlabel(\"x [m]\")\n",
    "ax1.set_title(\"Layer 1\")\n",
    "for ax in [ax0, ax1]:\n",
    "    ax.set_ylabel(\"head [m]\")\n",
    "    ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infiltration between two rivers\n",
    "\n",
    "Problem from Ch. 5 Analytical Groundwater Modeling: Theory and Applications Using Python (Bakker & Post, 2022)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = [20.0]\n",
    "H = 10.0\n",
    "Saq = [0.1]\n",
    "L = 1000.0\n",
    "N = 1e-3\n",
    "z = [0, -H]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlconf = tft.ModelXsection(naq=1, tmin=1, tmax=1e3)\n",
    "\n",
    "left = tft.XsectionMaq(\n",
    "    mlconf,\n",
    "    -np.inf,\n",
    "    -L / 2,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    # c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "inf = tft.XsectionMaq(\n",
    "    mlconf,\n",
    "    -L / 2,\n",
    "    L / 2,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    # c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    "    tsandN=[(0.0, N)],\n",
    ")\n",
    "right = tft.XsectionMaq(\n",
    "    mlconf,\n",
    "    L / 2,\n",
    "    np.inf,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    Saq=Saq,\n",
    "    # c=c,\n",
    "    topboundary=\"conf\",\n",
    "    phreatictop=True,\n",
    ")\n",
    "\n",
    "d = -1e-3\n",
    "hls_left = tft.River1D(mlconf, xls=-L / 2 + d, tsandh=[(0, 0.0)], layers=[0])\n",
    "hls_right = tft.River1D(mlconf, xls=L / 2 - d, tsandh=[(0, 0.0)], layers=[0])\n",
    "\n",
    "mlconf.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlss = tfs.ModelMaq(kaq=k, z=z, c=c, topboundary=\"conf\")\n",
    "\n",
    "inhom_left = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    -np.inf,\n",
    "    -L / 2,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "inhom_mid = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    -L / 2,\n",
    "    L / 2,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    "    N=N,\n",
    ")\n",
    "\n",
    "inhom_right = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    L / 2,\n",
    "    np.inf,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"conf\",\n",
    ")\n",
    "\n",
    "hls_left = tfs.River1D(mlss, xls=-L / 2 + d, hls=0.0, layers=[0])\n",
    "hls_right = tfs.River1D(mlss, xls=L / 2 - d, hls=0.0, layers=[0])\n",
    "\n",
    "mlss.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-L / 2 + 1e-5, L / 2 - 1e-5, 101)\n",
    "y = np.zeros_like(x)\n",
    "t = [1, 2, 3]\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(1, 2, sharex=True, sharey=False, figsize=(10, 3))\n",
    "\n",
    "for i in range(len(t)):\n",
    "    h = mlconf.headalongline(x, y, t[i])\n",
    "    ax0.plot(x, h[0].squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "\n",
    "t = [40, 80, 160]\n",
    "for i in range(len(t)):\n",
    "    h = mlconf.headalongline(x, y, t[i])\n",
    "    ax1.plot(x, h[0].squeeze(), label=f\"t={t[i]:.0f} d\", color=f\"C{i + 3}\")\n",
    "\n",
    "hss = mlss.headalongline(x, y)\n",
    "ax1.plot(x, hss[0].squeeze(), color=\"k\", ls=\"dashed\", label=\"steady\")\n",
    "\n",
    "ax0.legend(loc=(0, 1), frameon=False, ncol=3, fontsize=\"small\")\n",
    "ax1.legend(loc=(0, 1), frameon=False, ncol=4, fontsize=\"small\")\n",
    "ax0.set_ylabel(\"head [m]\")\n",
    "for ax in [ax0, ax1]:\n",
    "    ax.set_xlabel(\"x [m]\")\n",
    "    ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Change in waterlevel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = [10.0]\n",
    "H = 10.0\n",
    "z = [1.0, 0, -H]\n",
    "S = [1e-2]\n",
    "c = [100.0]\n",
    "delh = 2.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlsemi = tft.ModelXsection(naq=1, tmin=1, tmax=1e2)\n",
    "\n",
    "inhom_left = tft.XsectionMaq(\n",
    "    mlsemi,\n",
    "    -np.inf,\n",
    "    0.0,\n",
    "    kaq=k,\n",
    "    Saq=Saq,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"semi\",\n",
    "    phreatictop=False,\n",
    "    tsandhstar=[(0.0, delh)],\n",
    ")\n",
    "inhom_right = tft.XsectionMaq(\n",
    "    mlsemi,\n",
    "    0.0,\n",
    "    np.inf,\n",
    "    kaq=k,\n",
    "    Saq=Saq,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"semi\",\n",
    "    phreatictop=False,\n",
    "    tsandhstar=[(0.0, 0.0)],\n",
    ")\n",
    "\n",
    "# hls_right = tft.River1D(mlsemi, xls=300, tsandh=[(0, 0.0)], layers=[0])\n",
    "\n",
    "mlsemi.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlss = tfs.ModelMaq(\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=0.0,\n",
    ")\n",
    "\n",
    "inhom_left_ss = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    -np.inf,\n",
    "    0.0,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=delh,\n",
    ")\n",
    "inhom_right_ss = tfs.XsectionMaq(\n",
    "    mlss,\n",
    "    0.0,\n",
    "    np.inf,\n",
    "    kaq=k,\n",
    "    z=z,\n",
    "    c=c,\n",
    "    topboundary=\"semi\",\n",
    "    hstar=0.0,\n",
    ")\n",
    "mlss.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lab = np.sqrt(k[0] * H * c[0])\n",
    "L = 2 * 3 * lab\n",
    "ax = inhom_left.plot(params=True)\n",
    "ax = inhom_right.plot(ax=ax, params=True)\n",
    "mlsemi.elementlist[0].plot(ax=ax)\n",
    "ax.set_xlim(-100, 100);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-L / 2, L / 2, 101)\n",
    "y = np.zeros_like(x)\n",
    "t = [1, 10, 100]\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True, figsize=(10, 4))\n",
    "\n",
    "h = mlsemi.headalongline(x, y, t)\n",
    "q = np.zeros_like(h)\n",
    "for i in range(len(x)):\n",
    "    q[:, :, i] = mlsemi.disvec(x[i], y[i], t)[0]\n",
    "for i in range(len(t)):\n",
    "    ax0.plot(x, h[0, i].squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "    ax1.plot(x, q[0, i].squeeze(), label=f\"t={t[i]:.0f} d\")\n",
    "\n",
    "hss = mlss.headalongline(x, y)\n",
    "qss = mlss.disvecalongline(x, y)\n",
    "ax0.plot(x, hss[0].squeeze(), color=\"k\", ls=\"dashed\", label=\"steady\")\n",
    "ax1.plot(x, qss[0].squeeze(), color=\"k\", ls=\"dashed\", label=\"steady\")\n",
    "\n",
    "ax0.legend(loc=(0, 1), frameon=False, ncol=4, fontsize=\"small\")\n",
    "ax0.set_ylabel(\"head [m]\")\n",
    "ax1.set_xlabel(\"x [d]\")\n",
    "ax1.set_ylabel(\"q [m$^2$/d]\")\n",
    "for ax in [ax0, ax1]:\n",
    "    ax.grid()\n",
    "fig.align_ylabels()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}