{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test line-sink strings\n",
    "\n",
    "This notebook tests line-sink string elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import timml as tml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One `HeadLineSinkString` in one layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tml.ModelMaq(kaq=[10, 20, 5], z=[0, -20, -40, -80, -90, -140], c=[4000, 10000])\n",
    "rf = tml.Constant(ml, xr=0, yr=10000, hr=20, layer=0)\n",
    "ls1 = tml.HeadLineSinkString(\n",
    "    ml, xy=[(-100, 0), (0, 0), (100, 0), (100, 50)], hls=10, layers=0\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "for i in range(3):\n",
    "    plt.subplot(1, 3, i + 1)\n",
    "    ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=i, newfig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `HeadLineSinkString` in two layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tml.ModelMaq(kaq=[10, 20, 5], z=[0, -20, -40, -80, -90, -140], c=[4000, 10000])\n",
    "rf = tml.Constant(ml, xr=0, yr=10000, hr=20, layer=0)\n",
    "ls1 = tml.HeadLineSinkString(\n",
    "    ml, xy=[(-100, 0), (0, 0), (100, 0), (100, 50)], hls=10, layers=[0, 1]\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "for i in range(3):\n",
    "    plt.subplot(1, 3, i + 1)\n",
    "    ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=i, newfig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `HeadLineSinkString`, each section in different layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tml.ModelMaq(kaq=[10, 20, 5], z=[0, -20, -40, -80, -90, -140], c=[4000, 10000])\n",
    "rf = tml.Constant(ml, xr=0, yr=10000, hr=20, layer=0)\n",
    "ls1 = tml.HeadLineSinkString(\n",
    "    ml, xy=[(-100, 0), (0, 0), (100, 0), (100, 50)], hls=10, layers=[0, 1, 0]\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TODO: make sure segments are plotted correctly in each layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "for i in range(3):\n",
    "    plt.subplot(1, 3, i + 1)\n",
    "    ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=i, newfig=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=1, newfig=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A 2D array (NLS, NLAYERS)\n",
    "Layers can always be specified as a 2D array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = tml.ModelMaq(kaq=[10, 20, 5], z=[0, -20, -40, -80, -90, -140], c=[4000, 10000])\n",
    "rf = tml.Constant(ml, xr=0, yr=10000, hr=20, layer=0)\n",
    "ls1 = tml.HeadLineSinkString(\n",
    "    ml,\n",
    "    xy=[(-100, 0), (0, 0), (100, 0), (100, 50)],\n",
    "    hls=10,\n",
    "    layers=[[0, 1], [0, 1], [0, 1]],\n",
    ")\n",
    "ml.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "for i in range(3):\n",
    "    plt.subplot(1, 3, i + 1)\n",
    "    ml.plots.contour([-200, 200, -100, 100], ngr=50, layers=i, newfig=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}