{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nonuniform flow model\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{d}{dx} \\left( \\frac{\\beta Q^2}{2gA^2} + H \\right)= -i_e\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, parallel=False)\n",
    "def NonUniformflow(sections, Q, Hdb):\n",
    "    g = float(9.8)\n",
    "    dhini = float(0.5)\n",
    "    H = np.empty_like(Q)\n",
    "    H[0] = Hdb\n",
    "    arr = sections[0].calIeAlphaBetaRcUsubABS(Q[0], H[0])\n",
    "    ied = arr[0]\n",
    "    Ad = arr[-3]\n",
    "    \n",
    "    Hd = H[0]\n",
    "    Qd = Q[0]\n",
    "    \n",
    "    for i in range(1, len(Q)):\n",
    "        d = i - 1\n",
    "        sc, sd = sections[i], sections[d]\n",
    "        Qc = Q[i]\n",
    "        Hc = sc.calHcABS( Qc )[0]\n",
    "        arr = sc.calIeAlphaBetaRcUsubABS(Qc, Hc)\n",
    "        iec = arr[0]\n",
    "        Ac  = arr[-3]\n",
    "    \n",
    "        dx = sc.distance - sd.distance\n",
    "        \n",
    "        E1 = 0.5/g*Qc**2.0/Ac**2.0 + Hc\n",
    "        E2 = 0.5/g*Qd**2.0/Ad**2.0 + Hd + 0.5*dx*(ied + iec)\n",
    "        \n",
    "        if E2 < E1 :\n",
    "            H[i] = Hc\n",
    "        else : \n",
    "            Hc = Hc + float(0.001)\n",
    "            dh = dhini\n",
    "            for n in range(1000):\n",
    "                arr = sc.calIeAlphaBetaRcUsubABS(Qc, Hc)\n",
    "                iec = arr[0]\n",
    "                Ac  = arr[-3]\n",
    "    \n",
    "                E1 = 0.5/g*Qc**2.0/Ac**2.0 + Hc\n",
    "                E2 = 0.5/g*Qd**2.0/Ad**2.0 + Hd + 0.5*dx*(ied + iec)\n",
    "                \n",
    "                if np.abs(E1 - E2) < 0.00001 : \n",
    "                    break\n",
    "                elif E1 > E2 :\n",
    "                    dh *= float(0.5)\n",
    "                    Hc -= dh\n",
    "                else:\n",
    "                    Hc += dh\n",
    "                \n",
    "            H[i] = Hc\n",
    "            \n",
    "        Qd, Hd, ied, Ad = Qc, Hc, iec, Ac\n",
    "                \n",
    "    return H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# unsteady flow model\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    &\\frac{\\partial A}{\\partial t} + \\frac{\\partial Q}{\\partial x} = 0 \\\\\n",
    "    &\\frac{\\partial Q}{\\partial t} + \\frac{\\partial }{\\partial x}\\left(\\dfrac{\\beta Q^2}{A}\\right) \n",
    "    + gA \\frac{\\partial H}{\\partial x} + gAi_e = 0\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True, parallel=False)\n",
    "def UnSteadyflow(sections, A, Q, H, Abound, Qbound, dt):\n",
    "    g = float(9.8)\n",
    "    imax = len(A)\n",
    "    Anew, Qnew, Hnew = np.zeros(imax), np.zeros(imax), np.zeros(imax)\n",
    "    ie = np.zeros(imax)\n",
    "    Beta = np.zeros(imax)\n",
    "    \n",
    "# continuous equation\n",
    "    for i in range(1, imax-1) : \n",
    "        dx = 0.5*(sections[i-1].distance - sections[i+1].distance)\n",
    "        Anew[i] = A[i] - dt * ( Q[i] - Q[i-1] ) / dx\n",
    "        \n",
    "    Anew[imax-1] = Abound\n",
    "    Anew[0] = Anew[1]\n",
    "#     Anew[0] = (Anew[1] - A[1]) + A[0]\n",
    "    \n",
    "    for i in range(imax) : \n",
    "        s = sections[i]\n",
    "        Hnew[i], _, _ = s.A2HBS(Anew[i], H[i])\n",
    "        arr = s.calIeAlphaBetaRcUsubABS(Q[i], H[i])\n",
    "        ie[i] = arr[0]\n",
    "        Beta[i] = arr[2]\n",
    "    \n",
    "# moumentum equation\n",
    "    for i in range(1, imax-1): \n",
    "        ic, im, ip = i, i-1, i+1\n",
    "        dxp = sections[ic].distance - sections[ip].distance\n",
    "        dxm = sections[im].distance - sections[ic].distance\n",
    "        dxc = 0.5*(sections[im].distance - sections[ip].distance)\n",
    "        \n",
    "        Cr1 = 0.5*( Q[ic]/A[ic] + Q[ip]/A[ip] )*dt/dxp\n",
    "        Cr2 = 0.5*( Q[ic]/A[ic] + Q[im]/A[im] )*dt/dxm\n",
    "        dHdx1 = ( Hnew[ip] - Hnew[ic] ) / dxp\n",
    "        dHdx2 = ( Hnew[ic] - Hnew[im] ) / dxm\n",
    "        dHdx = (float(1.0) - Cr1) * dHdx1 + Cr2 * dHdx2\n",
    "        \n",
    "        Qnew[ic] = Q[ic] - dt * ( Beta[ic]*Q[ic]**2/A[ic] - Beta[im]*Q[im]**2/A[im] ) / dxc \\\n",
    "                         - dt * g * Anew[ic] * dHdx \\\n",
    "                         - dt * g * A[ic] * ie[ic] \n",
    "        \n",
    "    Qnew[imax-1] = Qnew[imax-2]\n",
    "    Qnew[0] = Qbound\n",
    "        \n",
    "    return Anew, Qnew, Hnew"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}