{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# License\n",
    "https://raw.githubusercontent.com/computational-sediment-hyd/a-rudimentary-knowledge-of-river-bed-variation/master/LICENSE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# import module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import numba\n",
    "import sys\n",
    "import os\n",
    "import json\n",
    "import glob as glob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Governing Equation of River flow "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}\n",
    "    \\frac{\\partial A}{\\partial t}+\\frac{\\partial Q}{\\partial x} = 0 \\\\\n",
    "\\end{align}\n",
    "\n",
    "\\begin{align}\n",
    "    \\frac{\\partial Q}{\\partial t}+\\frac{\\partial}{\\partial x}\\left(\\frac{ Q^2}{ A}\\right) + gA\\frac{\\partial H}{\\partial x}+gAi_e = 0\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def flowmodel(A, Q, Adown, Qup, dAb, zb, B, dx, dt, g, manning):\n",
    "    imax = len(A)\n",
    "    Anew, Qnew = np.zeros(imax), np.zeros(imax)\n",
    "    \n",
    "# continuous equation\n",
    "    for i in numba.prange(1, imax-1) : Anew[i] = A[i] - dt * ( Q[i] - Q[i-1] ) / dx\n",
    "        \n",
    "    Anew[0], Anew[-1] = Anew[1], Adown\n",
    "    \n",
    "# moumentum equation\n",
    "    for i in numba.prange(1,imax-1): \n",
    "        ip, ic, im = (i+1, i, i-1) \n",
    "        Cr1 = 0.5*( Q[ic]/A[ic] + Q[ip]/A[ip] )*dt/dx\n",
    "        Cr2 = 0.5*( Q[ic]/A[ic] + Q[im]/A[im] )*dt/dx\n",
    "        dHdx1 = ( Anew[ip]/B[ip] + zb[ip] + dAb[ip]/B[ip] - Anew[ic]/B[ic] - zb[ic] - dAb[ic]/B[ic] ) / dx\n",
    "        dHdx2 = ( Anew[ic]/B[ic] + zb[ic] + dAb[ic]/B[ic] - Anew[im]/B[im] - zb[im] - dAb[im]/B[im] ) / dx\n",
    "        dHdx = (1.0 - Cr1) * dHdx1 + Cr2 * dHdx2\n",
    "        \n",
    "        Qnew[ic] = Q[ic] - dt * ( Q[ic]**2/A[ic] - Q[im]**2/A[im] ) / dx  \\\n",
    "                   - dt * g * Anew[ic] * dHdx \\\n",
    "                   - dt * g * A[ic] * manning**2 * Q[ic]**2 / B[ic]**2 / ( A[ic]/B[ic] )**(10.0/3.0)\n",
    "                \n",
    "    Qnew[0], Qnew[-1] = Qup, Qnew[-2]\n",
    "            \n",
    "    return Anew, Qnew "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Governing Equation of River Bed : Hirano model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\\begin{align}\n",
    "    (1-\\lambda)\\frac{\\partial A_{b}}{\\partial t}+\\frac{\\partial }{\\partial x} \\sum_{i=1}^n ( Q_{bi}P_i)= 0 \\\\\n",
    "\\end{align}\n",
    "\n",
    "\\begin{align}\n",
    "    \\frac{\\partial P_i}{\\partial t} &= - \\frac{1}{(1-\\lambda)E_dB} \\frac{\\partial (Q_{bi}P_i)}{\\partial x} - \\frac{P_{si}}{E_d B}\\frac{\\partial A_b}{\\partial t} \\\\\n",
    "    &= - \\frac{1}{E_d B}\\left(\\frac{\\partial A_{bi}}{\\partial t} + P_{si}\\frac{\\partial A_b}{\\partial t}\\right) \n",
    "\\end{align}\n",
    "\n",
    "\\begin{align}\n",
    "&P_{si} = \n",
    "\\left\\{ \\begin{array}{ll}\n",
    "    P_i \\mbox{ in exchange layer} & \\left(\\dfrac{\\partial A_b}{\\partial t} > 0 \\right) \\\\\n",
    "    P_i \\mbox{ under exchange layer} & \\left(\\dfrac{\\partial A_b}{\\partial t} < 0 \\right) \\\\\n",
    "\\end{array} \\right. \\\\\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "$Q_{bi}$ : Ashida-Michiue Eq. and  Egiazaroff Eq.\n",
    "\n",
    "$E_d$ : thickness of exchange layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}\n",
    "    \\frac{\\tau_{*c i}}{\\tau_{*cm}} &= 0.85 \\frac{D_m}{D_i} \\qquad & \\left(\\frac{D_i}{D_m} < 0.4 \\right) \\\\\n",
    "    \\dfrac{\\tau_{*ci}}{\\tau_{*cm}} &= \\left( \\dfrac{ \\displaystyle \\log_e 19 }{ \\displaystyle \\log_e \\left(19\\dfrac{D_i}{D_m} \\right) } \\right)^2 \\qquad & \\left(\\frac{D_i}{D_m} \\geq 0.4 \\right)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def sedimentmodel(dAb, dratio, A, Q, B, dx, dt, g, manning, dsize, dratioStandard, hExlayer, Qbup):\n",
    "    rhosw = 1.65 # grain specific gravity in water\n",
    "    porosity = 0.4\n",
    "    tscAve = 0.05 # critical tractive force of average grain size\n",
    "    \n",
    "    dAbnew = np.zeros_like( dAb )\n",
    "    drationew = np.zeros_like( dratio )\n",
    "    Qbsub = np.zeros_like( dratio )\n",
    "    dAbsub = np.zeros_like( dratio )\n",
    "    \n",
    "    imax, lmax = len(dratio), len(dratio[0])\n",
    "    \n",
    "    for i in numba.prange(imax):\n",
    "        Ap, Qp, dr, Bp = A[i], Q[i], dratio[i], B[i]\n",
    "        dAve = np.sum(dsize * dr)\n",
    "        us = np.sqrt(g * Ap/Bp) * Qp/Ap * manning / (Ap/Bp)**(2/3)\n",
    "        tsAve = us**2.0/rhosw/g/dAve\n",
    "        use = Qp/Ap / ( 6.0 + 2.5 * np.log( Ap/Bp/dAve/( 1.0+2.0*tsAve) ) )\n",
    "        Kc=1.0\n",
    "        \n",
    "        for l in range(lmax):\n",
    "            dri, dsi = dr[l], dsize[l]\n",
    "            ts  = us**2.0 /rhosw/g/dsi\n",
    "            tse = use**2.0/rhosw/g/dsi\n",
    "            \n",
    "# Egiazaroff Eq. \n",
    "            x = dsi/dAve\n",
    "            tscbytscm = (np.log(19)/np.log(19*x))**2 if x > 0.4 else 0.85/x\n",
    "            tsc = Kc * tscbytscm * tscAve \n",
    "       \n",
    "            if ts > tsc :\n",
    "# Ashida-Michiue Eq. \n",
    "                Qbsub[i][l] = np.sqrt(rhosw * g * dsi**3.0) \\\n",
    "                        * 17.0 * tse**1.5 * ( 1.0 - tsc / ts) * ( 1.0 - np.sqrt(tsc / ts) ) \\\n",
    "                        * dri * B[i]\n",
    "            else:\n",
    "                Qbsub[i][l] = 0.0\n",
    "                \n",
    "    for i in numba.prange(imax):\n",
    "        Qbs = Qbsub[i]\n",
    "        if i == 0:\n",
    "            if np.min(Qbup) < 0.0 : # equilibrium sand supply\n",
    "                dAbsub[i][:] = 0.0\n",
    "            else:\n",
    "                for l in range(lmax):\n",
    "                    dAbsub[i][l] = - dt / (1.0 - porosity) * ( Qbsub[i][l]-Qbup[l] )/dx\n",
    "                    if np.abs(dAbsub[i][l]) > hExlayer :\n",
    "                        print(i) ; print('dzsub-error')\n",
    "        else:\n",
    "            for l in range(lmax):\n",
    "                dAbsub[i][l] = - dt / (1.0 - porosity) * ( Qbsub[i][l] - Qbsub[i-1][l] )/dx\n",
    "                if np.abs(dAbsub[i][l]) > hExlayer : \n",
    "                    print(i) ; print('dzsub-error')\n",
    "                \n",
    "    for i in numba.prange(imax):\n",
    "        Qbs = Qbsub[i]\n",
    "# update dAb \n",
    "        dAball = np.sum( dAbsub[i] )\n",
    "        if np.abs(dAball) > hExlayer:\n",
    "            print(i) ; print('dz-error')\n",
    "            \n",
    "        dAbnew[i] = dAb[i] + dAball\n",
    "        \n",
    "# update dratio\n",
    "        dratioIn = dratioStandard[i][:] if dAball < 0.0 else dratio[i][:]\n",
    "            \n",
    "        for l in range(lmax):\n",
    "            drationew[i][l] = dratio[i][l] + ( dAbsub[i][l] - dAball * dratioIn[l] ) /hExlayer/B[i]\n",
    "            if drationew[i][l] < 0.0 : drationew[i][l] = 0.0\n",
    "            \n",
    "# correct ration so that sum of ration become 100% \n",
    "        sumdr = np.sum( drationew[i] )\n",
    "        drationew[i] /= sumdr\n",
    "        \n",
    "    return dAbnew, drationew, Qbsub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def bedvariation(\n",
    "dx,dt,manning,totalTime,outTimeStep,RunUpTime\n",
    ",dsize ,dratioStandard ,dratio ,hExlayer ,A ,Q ,B ,zb ,dAb ,Qup ,Adown \n",
    ",outputfilename, screenclass\n",
    "):\n",
    "    \n",
    "    g = 9.8\n",
    "    Qbup = np.full( ( len(dsize) ), -9999.0 ) # when equilibrium sand supply, set qbup to minus value\n",
    "        \n",
    "# run-up calculation\n",
    "    for i in range(int(RunUpTime/dt)):\n",
    "        ib = ( ( dAb[-2]/B[-2] + zb[-2] ) - ( dAb[-1]/B[-1] + zb[-1] ) )/dx\n",
    "        Adownp = Adown(0.0, Q[-1], dAb[-1]/B[-1], ib)\n",
    "        Qupp = Qup(0.0)\n",
    "        A, Q = flowmodel(A, Q, Adownp, Qupp, dAb, zb, B, dx, dt, g, manning)\n",
    "        \n",
    "# main calculation\n",
    "    for i in range( int(totalTime/dt) ):\n",
    "        \n",
    "# cal bed variation\n",
    "        dAb, dratio, Qbsub = sedimentmodel(dAb, dratio, A, Q, B, dx, dt, g, manning, dsize, dratioStandard, hExlayer, Qbup)\n",
    "        \n",
    "# cal water profile\n",
    "        ib = ( ( dAb[-2]/B[-2] + zb[-2] ) - ( dAb[-1]/B[-1] + zb[-1] ) )/dx\n",
    "        Adownp = Adown(i*dt, Q[-1], dAb[-1]/B[-1], ib)\n",
    "        Qupp = Qup(i*dt)\n",
    "        A, Q = flowmodel(A, Q, Adownp, Qupp, dAb, zb, B, dx, dt, g, manning)\n",
    "        \n",
    "# output \n",
    "        if( int(i*dt) % int(outTimeStep) ) == 0 :\n",
    "            print( str( int(i*dt) ) + ' second') \n",
    "            profile = []\n",
    "            for ii, (ap, qp, z, r, q) in enumerate(zip(A, Q, dAb, dratio, Qbsub)) : \n",
    "                profile.append({'distance' : int(ii*dx), 'A':ap, 'Q':qp, 'dAb':z, 'ratio':list(r), 'Qb':list(q)})\n",
    "            json.dump( {'time':int(i*dt), 'profile' : profile}, open('%010d' % int(i*dt) + 'sec.json', 'w') )\n",
    "            \n",
    "# join json files\n",
    "    d = [ json.load( open(f, 'r') ) for f in glob.glob('*sec.json') ]\n",
    "# delete json files\n",
    "    for f in glob.glob('*sec.json') : os.remove(f)\n",
    "        \n",
    "    d.sort(key=lambda x: x['time'])\n",
    "    \n",
    "    cond = {'width':list(B), 'elevation':list(zb), 'screenclass':list(screenclass), 'dsize':list(dsize) }\n",
    "    json.dump( {'condition':cond, 'output':d}, open(outputfilename, 'w') )"
   ]
  }
 ],
 "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.6.4"
  },
  "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": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "165px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}