{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def conEq(dep, qx, qy, dzb, dt, dx, dy, ibx, hmin, hbuf, hdown, periodic=True):\n",
    "    \n",
    "    imax, jmax = len(dep), len(dep[0])\n",
    "    depn = np.zeros_like(dep, dtype=np.float64)\n",
    "    fluxx = np.zeros((imax+1, jmax), dtype=np.float64)\n",
    "    fluxy = np.zeros((imax, jmax+1), dtype=np.float64)\n",
    "    modflux = np.full( (imax, jmax), False)\n",
    "    \n",
    "    gravity = float( 9.8 )\n",
    "    \n",
    "    f = 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",
    "    \n",
    "    def flux(Qp, Qm, depp, depm, zbp, zbm, ib, delta) : \n",
    "        r = f(Qp, Qm)\n",
    "        if ( (depm + zbm) <= (zbp + hbuf - ib*delta)) and (depp <= hbuf) : r = 0.0\n",
    "        if ( (depp + zbp) <= (zbm + hbuf + ib*delta)) and (depm <= hbuf) : r = 0.0\n",
    "            \n",
    "        return r\n",
    "        \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in range( jmax ):\n",
    "            c, xm = (i,j), (i-1,j)\n",
    "            fluxx[c] = flux(qx[c], qx[xm], dep[c], dep[xm], dzb[c], dzb[xm], ibx, dx)\n",
    "            \n",
    "    if periodic :\n",
    "# boundary : periodic\n",
    "        fluxx[-1,:] = fluxx[0,:] \n",
    "    else:\n",
    "        for j in numba.prange( jmax ): fluxx[-1,j] = fluxx[-2,j] # qx[-1,j] # if qx[-1,j] > 0.0 else qx[,j]\n",
    "# normal            \n",
    "#         for j in numba.prange( jmax ): fluxx[-1,j] = qx[-1,j] if qx[-1,j] > 0.0 else 0.0\n",
    "        \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in range( 1, jmax ):\n",
    "            c, ym = (i,j), (i,j-1)\n",
    "            fluxy[c] = flux(qy[c], qy[ym], dep[c], dep[ym], dzb[c], dzb[ym], 0.0, dy)\n",
    "            \n",
    "# wall boundary \n",
    "    fluxy[:,0] = 0.0 \n",
    "    fluxy[:,-1] = 0.0 \n",
    "    \n",
    "#     for i in numba.prange( imax ):\n",
    "#         fluxy[i,-1] = qy[i,-1] if qy[i,-1] > 0.0 else 0.0\n",
    "#         fluxy[i, 0] = qy[i, 0] if qy[i, 0] < 0.0 else 0.0    \n",
    "    \n",
    "    nis = 0 if periodic else 1\n",
    "# limiter --------------------------------------------------------------\n",
    "# 水深が負になる際に質量保存を満たすためにフラックスを修正する\n",
    "    for i in range(nis, imax):\n",
    "        for j in range(jmax):\n",
    "            if dep[c] > hmin :\n",
    "                c, xp, yp = (i, j), (i+1, j), (i, j+1)\n",
    "                fxp = fluxx[xp] if fluxx[xp] > 0.0 else 0.0\n",
    "                fxm = -fluxx[c] if fluxx[c] < 0.0 else 0.0\n",
    "                fyp = fluxy[yp] if fluxy[yp] > 0.0 else 0.0\n",
    "                fym = -fluxy[c] if fluxy[c] < 0.0 else 0.0\n",
    "                V =  dep[c]*dx*dy - hmin*dx*dy \n",
    "                Vq = ( fxp*dy + fxm*dy + fyp*dx + fym*dx )*dt\n",
    "                if V <  Vq:\n",
    "                    alfa = V / Vq - 0.001\n",
    "                    if fluxx[xp] > 0.0 : fluxx[xp] *= alfa \n",
    "                    if fluxx[c]  < 0.0 : fluxx[c]  *= alfa\n",
    "                    if fluxy[yp] > 0.0 : fluxy[yp] *= alfa\n",
    "                    if fluxy[c]  < 0.0 : fluxy[c]  *= alfa\n",
    "                        \n",
    "                    modflux[c] = True\n",
    "# ------------------------------------------------------------------------\n",
    "    n = 0\n",
    "    for i in numba.prange(nis, imax):\n",
    "        for j in range(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 : \n",
    "                n += 1\n",
    "#                 print('dep-error')\n",
    "#                 print( modflux[c], depn[c], fluxx[xp], fluxx[c], fluxy[yp], fluxy[c] )\n",
    "                fxp = fluxx[xp] if fluxx[xp] > 0.0 else 0.0\n",
    "                fxm = -fluxx[c] if fluxx[c] < 0.0 else 0.0\n",
    "                fyp = fluxy[yp] if fluxy[yp] > 0.0 else 0.0\n",
    "                fym = -fluxy[c] if fluxy[c] < 0.0 else 0.0\n",
    "                V =  dep[c]*dx*dy - hmin*dx*dy \n",
    "                Vq = ( fxp*dy + fxm*dy + fyp*dx + fym*dx )*dt\n",
    "#                 print(V,Vq)\n",
    "                \n",
    "                depn[c] = hmin\n",
    "                \n",
    "# upstream boundary                \n",
    "#     if periodic == False: depn[0][:] = depn[1][:]\n",
    "    if periodic == False: \n",
    "#         depn[0][:] = dep[0][:]\n",
    "#         for j in Qind : depn[0,j] = depn[1,j]  \n",
    "            \n",
    "        depn[0,:] = depn[1,:]  \n",
    "    \n",
    "# downstream boundary \n",
    "#     depn[-1][:] = hdown\n",
    "        depn[-1][:] = depn[-2][:]\n",
    "    \n",
    "    return depn, n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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  \\left(\\frac{\\partial^2 uh}{\\partial x^2}+\\frac{\\partial^2 uh}{\\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  \\left(\\frac{\\partial^2 vh}{\\partial x^2}+\\frac{\\partial^2 vh}{\\partial y^2} \\right)\n",
    "        = 0\n",
    "    \\end{align}\n",
    "$$ -->\n",
    "\n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "        &\\mathrm{case1} : \\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",
    "        - \\frac{\\partial }{\\partial x} \\left(h \\nu_t \\dfrac{\\partial u}{\\partial x}\\right)\n",
    "        - \\frac{\\partial }{\\partial y} \\left(h \\nu_t \\dfrac{\\partial u}{\\partial y}\\right)\n",
    "         = 0 \\\\\n",
    "        &\\mathrm{case2} : \\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",
    "        - \\frac{\\partial }{\\partial x} \\left(2h \\nu_t \\dfrac{\\partial u}{\\partial x}- \\dfrac{2}{3}hk \\right )\n",
    "        - \\frac{\\partial }{\\partial y}\\left( h \\nu_t \\left(\\dfrac{\\partial u}{\\partial y} + \\dfrac{\\partial v}{\\partial x}\\right) \\right)\n",
    "         = 0\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=True)\n",
    "def momentEq(dep, qx, qy, depn, dzb, dt, dx, dy, ibx, hmin, hbuf, direction, Qup, cManning, periodic=True):\n",
    "    #direction = 1:x, 2:y\n",
    "    gravity = float( 9.8 )\n",
    "\n",
    "    q = qx.copy() if direction == 1 else qy.copy()\n",
    "    u, v = qx/dep, qy/dep\n",
    "    Vdir = q/dep  \n",
    "    \n",
    "    Cf = gravity*cManning**2.0/dep**(1.0/3.0) \n",
    "    Vnorm = np.sqrt(u**2.0+v**2.0) \n",
    "    kenergy = 2.07*Cf*Vnorm**2\n",
    "    nut = 0.4/6.0*dep*np.sqrt(Cf)*np.abs(Vnorm)\n",
    "        \n",
    "    imax, jmax = len(q), len(q[0])\n",
    "    \n",
    "    qn = np.zeros_like(q, dtype=np.float64)\n",
    "    fluxx = np.zeros((imax+1, jmax), dtype=np.float64)\n",
    "    fluxy = np.zeros((imax, jmax+1), dtype=np.float64)\n",
    "    \n",
    "    tauxy = np.zeros_like(q, dtype=np.float64)\n",
    "            \n",
    "    f = lambda vp,vm,qp,qm : vm*qm if vp >= 0.0 and vm >= 0.0 else \\\n",
    "        (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",
    "    def flux1(vp, vm, qp, qm, depp, depm, zbp, zbm, ib, delta) : \n",
    "        r = f(vp,vm,qp,qm)\n",
    "        if ( (depm + zbm) <= (zbp + hbuf - ib*delta)) and (depp <= hbuf) : r = 0.0\n",
    "        if ( (depp + zbp) <= (zbm + hbuf + ib*delta)) and (depm <= hbuf) : r = 0.0\n",
    "            \n",
    "        return r\n",
    "    \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in range( jmax ):\n",
    "            c, xm = (i,j), (i-1,j)\n",
    "            fluxx[c] = flux1( u[c], u[xm], q[c], q[xm], dep[c], dep[xm], dzb[c], dzb[xm], ibx, dx )\n",
    "            \n",
    "# boundary : periodic\n",
    "    if periodic :\n",
    "        fluxx[-1,:] = fluxx[0,:] \n",
    "    else:\n",
    "        fluxx[0,:]  = -9999\n",
    "        for j in numba.prange( jmax ):\n",
    "            fluxx[-1,j] = u[-1,j]*q[-1,j] if u[-1,j] > 0.0 else 0.0 \n",
    "    \n",
    "    for i in numba.prange( imax ):\n",
    "        for j in range( 1, jmax ):\n",
    "            c, ym = (i,j), (i,j-1)\n",
    "            fluxy[c] = flux1( v[c], v[ym], q[c], q[ym], dep[c], dep[ym], dzb[c], dzb[ym], 0.0, dy )\n",
    "            \n",
    "# wall boundary \n",
    "    fluxy[:,-1] = float(0.0)\n",
    "    fluxy[:, 0] = float(0.0)\n",
    "#     for i in numba.prange( imax ):\n",
    "#         fluxy[i,-1] = v[i,-1]*q[i,-1] if v[i,-1] > 0.0 else 0.0\n",
    "#         fluxy[i, 0] = v[i, 0]*q[i, 0] if v[i,0 ] < 0.0 else 0.0\n",
    "    \n",
    "# tauxy   \n",
    "    for i in numba.prange(imax):\n",
    "        for j in range(jmax):\n",
    "            c = (i, j)\n",
    "            if periodic :\n",
    "                xp = (0, j) if i == imax-1 else (i+1, j)                \n",
    "            else:\n",
    "                xp = (i+1, j)                \n",
    "                \n",
    "            xm = (i-1, j)\n",
    "            yp = (i, j+1)\n",
    "            ym = (i, j-1)\n",
    "            \n",
    "            if j==0 :\n",
    "                tauxy[i,j] = nut[c]*( (u[yp]-u[c])/1.0/dy + (v[xp]-v[xm])/2.0/dx )\n",
    "            elif j==jmax-1 :\n",
    "                tauxy[i,j] = nut[c]*( (u[c]-u[ym])/1.0/dy + (v[xp]-v[xm])/2.0/dx )\n",
    "            else:\n",
    "                tauxy[i,j] = nut[c]*( (u[yp]-u[ym])/2.0/dy + (v[xp]-v[xm])/2.0/dx )\n",
    "    \n",
    "    nis = 0 if periodic else 1\n",
    "    \n",
    "    for i in numba.prange(nis, imax):\n",
    "        for j in range(jmax):    \n",
    "            c = (i, j)\n",
    "            if periodic :\n",
    "#                 xp = (0, j) if i == imax-1 else (i+1, j)                \n",
    "                if i == imax-1:\n",
    "                    xp = (0, j) \n",
    "                else:\n",
    "                    xp = (i+1, j)                \n",
    "            else:\n",
    "                xp = (i+1, j)                \n",
    "                \n",
    "            xm = (i-1, j)\n",
    "            yp = (i, j+1)\n",
    "            ym = (i, j-1)\n",
    "            \n",
    "            if depn[c] <= hbuf :\n",
    "                qn[c] = 0.0\n",
    "            else:\n",
    "            # pressure & gravity term\n",
    "                if (direction == 2) and (j == 0)  : \n",
    "                    dHdx =((depn[yp]+dzb[yp]) - (depn[c]+dzb[c]))/dy\n",
    "                elif (direction == 2) and (j == jmax-1)  : \n",
    "                    dHdx =((depn[c]+dzb[c]) - (depn[ym]+dzb[ym]))/dy\n",
    "#                 if direction == 2 and ((j == 0) or (j == jmax-1)) : \n",
    "#                     dHdx = 0.0\n",
    "#                 elif direction == 1 and i == imax-1 :\n",
    "#                     dHdx =((depn[i,j]+dzb[i,j]) - (depn[i-1,j]+dzb[i-1,j]))/dx\n",
    "                else :\n",
    "                    if direction == 1 :\n",
    "                        dp, dm, delta  = xp, xm, dx\n",
    "                        dib = ibx\n",
    "                    else :\n",
    "                        dp, dm, delta = yp, ym, dy\n",
    "                        dib = float(0.0)\n",
    "                        \n",
    "                    Vc, Vp, Vm = q[c]/dep[c], q[dp]/dep[dp], q[dm]/dep[dm]\n",
    "                    Hc, Hp, Hm = depn[c]+dzb[c], depn[dp]+dzb[dp], depn[dm]+dzb[dm]\n",
    "\n",
    "                    if(Hc <= dzb[dp] + hbuf - dib*delta) and depn[dp] <= hbuf :\n",
    "                        if(Hc <= dzb[dm] + hbuf + dib*delta) and depn[dm] <= hbuf :\n",
    "                            dHdx = 0.0\n",
    "                        else:\n",
    "                            dHdx = (Hc-Hm)/delta-dib\n",
    "                    elif(Hc <= dzb[dm] + hbuf + dib*delta) and depn[dm] <= hbuf :\n",
    "                        dHdx = (Hp-Hc)/delta-dib\n",
    "                    else :\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-dib, (Hc-Hm)/delta-dib\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-dib, (Hp-Hc)/delta-dib          \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 - dib\n",
    "                            dkhdx1 = dkhdx2 = ( 0.5*(khp+khc) - 0.5*(khc+khm) )/delta\n",
    "            \n",
    "                        w1, w2 = 1-Cr1**0.5, Cr2**0.5\n",
    "                        dHdx = w1 * dHdx1 + w2 * dHdx2   \n",
    "                \n",
    "# viscous sublayer\n",
    "#                 Vis = Cf[c] * Vnorm[c] * u[c] if direction == 1 else  Cf[c] * Vnorm[c] * v[c]\n",
    "                Vis = Cf[c] * Vnorm[c] * Vdir[c] \n",
    "# turbulence\n",
    "# side boundary : non-slip condition\n",
    "#                 if (i == imax-1) or (j == 0) or (i == jmax-1): \n",
    "                if (j == 0) or (j == jmax-1): \n",
    "                    turb = 0.0\n",
    "                else:\n",
    "                    \n",
    "                # case 1\n",
    "#                     turb = (0.5*(nut[xp] + nut[c]) * 0.5*(dep[xp] + dep[c]) * (Vdir[xp] - Vdir[c] )/dx \\\n",
    "#                          -  0.5*(nut[xm] + nut[c]) * 0.5*(dep[xm] + dep[c]) * (Vdir[c]  - Vdir[xm])/dx)/dx \\\n",
    "#                          + (0.5*(nut[yp] + nut[c]) * 0.5*(dep[yp] + dep[c]) * (Vdir[yp] - Vdir[c] )/dy \\\n",
    "#                          -  0.5*(nut[ym] + nut[c]) * 0.5*(dep[ym] + dep[c]) * (Vdir[c]  - Vdir[ym])/dy)/dy\n",
    "                    \n",
    "                # case 2\n",
    "                    nutpave = 0.5*nut[dp] + 0.5*nut[c]\n",
    "                    deppave = 0.5*dep[dp] + 0.5*dep[c]\n",
    "                    kpave = 0.5*kenergy[dp] + 0.5*kenergy[c]\n",
    "                    turb1p = deppave * 2.0*nutpave*(Vdir[dp]-Vdir[c])/delta - 2.0/3.0*kpave\n",
    "                        \n",
    "                    nutmave = 0.5*nut[dm] + 0.5*nut[c]\n",
    "                    depmave = 0.5*dep[dm] + 0.5*dep[c]\n",
    "                    kmave = 0.5*kenergy[dm] + 0.5*kenergy[c]\n",
    "                    turb1m = depmave * 2.0*nutmave*(Vdir[c]-Vdir[dm])/delta - 2.0/3.0*kmave\n",
    "                    \n",
    "                    turb1 = (turb1p - turb1m)/delta\n",
    "                    \n",
    "                    if direction == 1 :\n",
    "                        turb2 = (dep[yp]*tauxy[yp] - dep[ym]*tauxy[ym])/2.0/dy\n",
    "                    else:\n",
    "                        turb2 = (dep[xp]*tauxy[xp] - dep[xm]*tauxy[xm])/2.0/dx\n",
    "\n",
    "                    turb = turb1 + turb2 \n",
    "                    \n",
    "                sourcet = Vis - turb\n",
    "                \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 * sourcet\n",
    "                \n",
    "    if periodic == False :\n",
    "# upstream boundary\n",
    "        if direction == 2 : \n",
    "            qn[0,:] = 0.0\n",
    "        else :\n",
    "            updep = np.zeros_like(depn[0,:])\n",
    "            for j in range(jmax):\n",
    "                if depn[0,j] > hmin : updep[j] = depn[0,j]\n",
    "                    \n",
    "            alpha = Qup / dy / np.sum( updep[:]**(5/3) )\n",
    "            qn[0,:] = alpha * updep[:]**(5/3)\n",
    "            \n",
    "# downstream boundary\n",
    "#     qn[-1,:] = qn[-2,:]\n",
    "        \n",
    "    return qn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=False)\n",
    "def simulation(dep, qx, qy, dzb, dt, dx, dy, ibx, hmin, hbuf, Qup, hdown, cManning):\n",
    "    depn, count = conEq(dep, qx, qy, dzb, dt, dx, dy, ibx, hmin, hbuf, hdown, True)\n",
    "    qxn = momentEq(dep, qx, qy, depn, dzb, dt, dx, dy, ibx, hmin, hbuf, 1, Qup, cManning, True)\n",
    "    qyn = momentEq(dep, qx, qy, depn, dzb, dt, dx, dy, float(0), hmin, hbuf, 2, Qup, cManning, True)\n",
    "#     CFL = ((np.abs(qxn/depn) + np.sqrt(9.8*depn))/dx + ( np.abs(qyn/depn) + np.sqrt(9.8*depn) )/dy)*dt\n",
    "    CFL = ((np.abs(qxn/depn))/dx + ( np.abs(qyn/depn))/dy)*dt\n",
    "    CFLmax = np.max(CFL)\n",
    "    return depn, qxn, qyn, CFLmax, count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# main "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hmin = float(10.0**(-5)) \n",
    "hbuf = float(10.0**(-2)) \n",
    "\n",
    "dx, dy = float(1.0), float(1.0)\n",
    "\n",
    "nxmax, nymax = 1800, 90\n",
    "qx = np.zeros((nxmax,nymax), dtype=np.float64)\n",
    "qy = np.zeros_like(qx, dtype=np.float64)\n",
    "dep = np.full_like(qx, 6.0, dtype=np.float64)\n",
    "zb = np.zeros_like(qx, dtype=np.float64)\n",
    "cManning = np.full_like(qx, 0.025, dtype=np.float64)\n",
    "cManning[:,:30] = 0.1\n",
    "\n",
    "ibx = float(0.001)\n",
    "\n",
    "qx = ibx**0.5*dep**(float(5)/float(3))/cManning\n",
    "\n",
    "X = np.arange(0, dx*nxmax, dx) + 0.5*dx\n",
    "Y = np.arange(0, dy*nymax, dy) + 0.5*dy\n",
    "\n",
    "arr = np.random.rand(nxmax, nymax)/1000 + 1\n",
    "qx = qx * arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "t = float(0)\n",
    "dt = 0.04\n",
    "dtout= float(600.0)\n",
    "tmax = float(5.0*3600.+0.1)\n",
    "nout = 0\n",
    "\n",
    "dic = {}\n",
    "\n",
    "while tmax >= t :\n",
    "    t += dt\n",
    "    dep, qx, qy, CFLmax, count = simulation(dep, qx, qy, zb, dt, dx, dy, ibx, hmin, hbuf, float(0), float(0), cManning)\n",
    "    \n",
    "# update dt\n",
    "#     dt = np.round( dt * CFL/CFLmax, 5) \n",
    "    \n",
    "    if t >= nout*dtout :\n",
    "        print(t, dt, CFLmax, count)\n",
    "        dic['total_second'] = round(t, 2)\n",
    "        dss = xr.Dataset({'depth': (['x','y'], dep), 'u': (['x','y'], qx/dep), 'v': (['x','y'], qy/dep) }\n",
    "                          , coords={'x': X, 'y':Y}\n",
    "                          , attrs=dic )\n",
    "        \n",
    "        out = dss.to_netcdf('out' + str(nout).zfill(8) + '.nc')\n",
    "        out = dss.close()\n",
    "        del out\n",
    "        nout += 1 "
   ]
  }
 ],
 "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": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "167.767px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}