{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# import module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numba\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Continuous Equation\n",
    "\n",
    "$$\n",
    "    \\begin{align}\n",
    "        \\frac{\\partial h}{\\partial t}+\\frac{\\partial q_x}{\\partial x} +\\frac{\\partial q_y}{\\partial y} = 0\n",
    "    \\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def conEq(dep, qx, qy, dt, dx, dy, hmin):\n",
    "    imax, jmax = len(dep), len(dep[0])\n",
    "    depn = np.empty_like(dep)\n",
    "    fluxx = np.zeros((imax+1, jmax))\n",
    "    fluxy = np.zeros((imax, jmax+1))\n",
    "    \n",
    "#     flux = lambda Qp, Qm : Qm if Qp > 0.0 and Qm > 0.0 else (Qp if Qp < 0.0 and Qm < 0.0 else 0.5*Qp+0.5*Qm )\n",
    "    def flux( Qp, Qm) : \n",
    "        if Qp > 0.0 and Qm > 0.0 :\n",
    "            r = Qm\n",
    "        elif Qp < 0.0 and Qm < 0.0 :\n",
    "            r = Qp\n",
    "        else :\n",
    "# treatment of dry-bed\n",
    "            if Qp == 0.0 and Qm == 0.0 :\n",
    "                r = 0.0\n",
    "            elif Qp == 0.0 :\n",
    "                r = Qm if Qm > 0.0 else 0.0\n",
    "            elif Qm == 0.0 :\n",
    "                r = Qp if Qp < 0.0 else 0.0\n",
    "            else :\n",
    "                r = 0.5*Qp+0.5*Qm \n",
    "                \n",
    "        return r\n",
    "        \n",
    "    \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in numba.prange( jmax ):\n",
    "            c, xm, ym = (i,j), (i-1,j), (i,j-1)\n",
    "            fluxx[c] = flux(qx[c], qx[xm])\n",
    "            fluxy[c] = flux(qy[c], qy[ym])\n",
    "            \n",
    "# continuous boundary \n",
    "    fluxx[-1,:] = fluxx[0,:]\n",
    "    fluxy[:,-1] = fluxy[:,0]\n",
    "    \n",
    "    for i in numba.prange(imax):\n",
    "        for j in numba.prange(jmax):\n",
    "            c, xp, yp = (i, j), (i+1, j), (i, j+1)\n",
    "            depn[c] = dep[c] - dt*(fluxx[xp] - fluxx[c])/dx - dt*(fluxy[yp] - fluxy[c])/dy\n",
    "            if depn[c] < hmin*2 : depn[c] = hmin\n",
    "\n",
    "    return depn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Momentum Equation\n",
    "\n",
    "$$\n",
    "    \\begin{align}\n",
    "        \\frac{\\partial q_x}{\\partial t}+\\frac{\\partial u q_x}{\\partial x}+\\frac{\\partial v q_x}{\\partial y}+gh\\frac{\\partial H}{\\partial x}+\\frac{\\tau_{0x}}{\\rho} \n",
    "        - \\nu_t h \\left(\\frac{\\partial^2 u}{\\partial x^2}+\\frac{\\partial^2 u}{\\partial y^2} \\right)= 0 \\\\\n",
    "        \\frac{\\partial q_y}{\\partial t}+\\frac{\\partial u q_y}{\\partial x}+\\frac{\\partial v q_y}{\\partial y}+gh\\frac{\\partial H}{\\partial y}+\\frac{\\tau_{0y}}{\\rho}- \\nu_t h \\left(\\frac{\\partial^2 v}{\\partial x^2}+\\frac{\\partial^2 v}{\\partial y^2} \\right)\n",
    "        = 0\n",
    "    \\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, direction):\n",
    "    #direction = 1:x, 2:y\n",
    "    gravity = 9.8\n",
    "    cManing = 0.02\n",
    "    ib = 1.0/670.0\n",
    "\n",
    "    q = qx if direction == 1 else qy\n",
    "    u, v = qx/dep, qy/dep\n",
    "        \n",
    "    imax, jmax = len(q), len(q[0])\n",
    "    \n",
    "    qn = np.empty_like(q)\n",
    "    fluxx = np.zeros((imax+1, jmax))\n",
    "    fluxy = np.zeros((imax, jmax+1))\n",
    "\n",
    "    flux = lambda vp,vm,qp,qm : vm*qm if vp > 0.0 and vm > 0.0 \\\n",
    "                 else (vp*qp if vp < 0.0 and vm < 0.0 else (0.5*vp+0.5*vm)*(0.5*qp+0.5*qm) )\n",
    "    \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in numba.prange( jmax ):\n",
    "            c, xm, ym = (i,j), (i-1,j), (i,j-1)\n",
    "            fluxx[c] = flux( u[c], u[xm], q[c], q[xm] )\n",
    "            fluxy[c] = flux( v[c], v[ym], q[c], q[ym] )\n",
    "            \n",
    "# continuous boundary \n",
    "    fluxx[-1,:] = fluxx[0,:]\n",
    "    fluxy[:,-1] = fluxy[:,0]\n",
    "            \n",
    "    for i in numba.prange(imax):\n",
    "        for j in numba.prange(jmax):    \n",
    "            c = (i, j)\n",
    "            if depn[c] <= hmin*2 :\n",
    "                qn[c] = 0.0\n",
    "            else:\n",
    "        # pressure & gravity term\n",
    "                c, xm, ym = (i,j), (i-1, j), (i, j-1)\n",
    "                xp = (0, j) if i == imax - 1 else (i+1, j  )\n",
    "                yp = (i, 0) if j == jmax - 1 else (i  , j+1)\n",
    "        \n",
    "                if direction == 1 :\n",
    "                    dp, dm, delta = xp, xm, dx \n",
    "                else :\n",
    "                    dp, dm, delta = yp, ym, dy\n",
    "                    \n",
    "                Vc, Vp, Vm = q[c]/dep[c], q[dp]/dep[dp], q[dm]/dep[dm]\n",
    "                Hc, Hp, Hm = depn[c]+zb[c], depn[dp]+zb[dp], depn[dm]+zb[dm]\n",
    "    \n",
    "                if Vc > 0.0 and Vp > 0.0 and Vm > 0.0: \n",
    "                    Cr1, Cr2 = 0.5*(abs(Vc)+abs(Vp))*dt/delta, 0.5*(abs(Vc)+abs(Vm))*dt/delta\n",
    "                    dHdx1, dHdx2 = (Hp-Hc)/delta-ib, (Hc-Hm)/delta-ib\n",
    "                elif Vc < 0.0 and Vp < 0.0 and Vm < 0.0:\n",
    "                    Cr1, Cr2 = 0.5*(abs(Vc)+abs(Vm))*dt/delta, 0.5*(abs(Vc)+abs(Vp))*dt/delta\n",
    "                    dHdx1, dHdx2 = (Hc-Hm)/delta-ib, (Hp-Hc)/delta-ib          \n",
    "                else:\n",
    "                    Cr1 = Cr2 = 0.5*(abs(0.5*(Vc+Vp))+abs(0.5*(Vc+Vm)))*dt/delta\n",
    "                    dHdx1 = dHdx2 = (0.5*(Hc+Hp) - 0.5*(Hc+Hm)) / delta - ib\n",
    "    \n",
    "                w1, w2 = 1-Cr1**0.5, Cr2**0.5\n",
    "                dHdx = w1 * dHdx1 + w2 * dHdx2   \n",
    "\n",
    "#viscous sublayer\n",
    "                Cf = gravity*cManing**2.0/dep[c]**(1.0/3.0) \n",
    "                Vnorm = np.sqrt(u[c]**2.0+v[c]**2.0) \n",
    "                Vis = Cf * Vnorm * u[c] if direction == 1 else  Cf * Vnorm * v[c]\n",
    "#turbulence\n",
    "#                kenergy = 2.07*Cf*Vnorm**2\n",
    "                nut = 0.4/6.0*dep[c]*np.sqrt(Cf)*np.abs(Vnorm)\n",
    "    \n",
    "                turb = nut * ( u[xp] - 2.0*u[c] + u[xm] )/ dx**2 \\\n",
    "                     + nut * ( u[yp] - 2.0*u[c] + u[ym] )/ dy**2\n",
    "\n",
    "            \n",
    "                c, xp, yp = (i,j), (i+1,j), (i,j+1)\n",
    "                qn[c] = q[c] - dt * ( fluxx[xp] - fluxx[c] ) / dx \\\n",
    "                             - dt * ( fluxy[yp] - fluxy[c] ) / dy \\\n",
    "                             - dt * gravity * depn[c] * dHdx \\\n",
    "                             - dt * Vis \\\n",
    "                             + dt * dep[c] * turb\n",
    "\n",
    "    return qn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def simulation(dep, qx, qy, zb, dt, dx, dy, hmin):\n",
    "    depn = conEq(dep, qx, qy, dt, dx, dy, hmin)\n",
    "    qxn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, 1)\n",
    "    qyn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, 2)\n",
    "    CFL = ((np.abs(qxn/depn))/dx + (np.abs(qyn/depn))/dy)*dt\n",
    "    CFLmin = np.max(CFL)\n",
    "    return depn, qxn, qyn, CFLmin     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Main routine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## read elevation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset('dz.nc')\n",
    "dza = ds['dz']\n",
    "dz = dza.values\n",
    "ds.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "200.0 0.12994759757305765\n",
      "400.0 0.05324376553275483\n",
      "600.0 0.13465034344958737\n",
      "800.0 0.08917062844146038\n",
      "1000.0 0.06319218033460254\n",
      "1200.0 0.11995618722432487\n",
      "1400.0 0.07224942603327707\n",
      "1600.0 0.12386222201873731\n",
      "1800.0 0.06899316890484393\n",
      "2000.0 0.170605601280248\n",
      "2200.0 0.06517755577401307\n",
      "2400.0 0.1353129352137529\n",
      "2600.0 0.0615231283218484\n",
      "2800.0 0.07240785399189434\n",
      "3000.0 0.07694049275220537\n",
      "3200.0 0.06135978323502305\n",
      "3400.0 0.08224366525350822\n",
      "3600.0 0.061249249226995206\n",
      "Wall time: 21.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "hmin = 10.0**(-3)\n",
    "nxmax = len(dz)\n",
    "nymax = len(dz[0])\n",
    "dx = 5.0\n",
    "dy = 5.0\n",
    "dt = 0.1\n",
    "dtout= 200.0\n",
    "tmax = 3600.0\n",
    "\n",
    "qx = np.zeros((nxmax,nymax))\n",
    "qy = np.zeros_like(qx)\n",
    "zb = dz #np.zeros_like(qx)\n",
    "\n",
    "# Initial Condition\n",
    "dep = np.full_like(zb, 0.4)\n",
    "\n",
    "def timeGenerater():\n",
    "    it, itout = 0, 1\n",
    "    isTiter, isOut = True, False\n",
    "    yield it*dt, isTiter, isOut\n",
    "\n",
    "    while True:\n",
    "        it += 1\n",
    "        if it*dt >= itout*dtout:\n",
    "            itout += 1\n",
    "            isOut = True\n",
    "        else:\n",
    "            isOut = False\n",
    "\n",
    "        if it*dt >= tmax: isTiter = False\n",
    "\n",
    "        yield it*dt, isTiter, isOut\n",
    "        \n",
    "tgen = timeGenerater() \n",
    "t, isTiter, isOut = next(tgen)\n",
    "tarr = [t]\n",
    "deparr = np.array([dep])\n",
    "uarr = np.array([qx/dep])\n",
    "varr = np.array([qy/dep])\n",
    "\n",
    "nout= 1\n",
    "while isTiter:\n",
    "    dep, qx, qy, CFLmin = simulation(dep, qx, qy, zb, dt, dx, dy, hmin)\n",
    "    t, isTiter, isOut = next(tgen)\n",
    "    if isOut:\n",
    "        tarr.append(round(t,2))\n",
    "        deparr = np.append(deparr, np.array([dep]), axis=0)\n",
    "        uarr = np.append(uarr, np.array([qx/dep]), axis=0)\n",
    "        varr = np.append(varr, np.array([qy/dep]), axis=0)\n",
    "        print(t,CFLmin)\n",
    "        nout += 1\n",
    "        \n",
    "dss = xr.Dataset({'depth': (['t','x','y'], deparr), 'u': (['t','x','y'], uarr), 'v': (['t','x','y'], varr)}, coords={'xc': (('x', 'y'), dza['xc'])\n",
    "                                             , 'yc': (('x', 'y'), dza['yc'])\n",
    "                                             , 't': tarr})\n",
    "\n",
    "dss.to_netcdf('output.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## check total volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1524.5068734878676, 1527.6000000000004)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum( deparr[-1] ), np.sum( deparr[0] )"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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": "384px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}