<jupyter_start><jupyter_text>import module<jupyter_code>import numpy as np
import numba
import xarray as xr<jupyter_output><empty_output><jupyter_text>Continuous Equation$$    \begin{align}        \frac{\partial h}{\partial t}+\frac{\partial q_x}{\partial x} +\frac{\partial q_y}{\partial y} = 0    \end{align}$$<jupyter_code>@numba.jit(nopython=True, parallel=False)
def conEq(dep, qx, qy, zb, dt, dx, dy, add, hmin, hbuf):
    imax, jmax = len(dep), len(dep[0])
    depn = np.zeros_like(dep, dtype=np.float64)
    fluxx = np.zeros((imax+1, jmax), dtype=np.float64)
    fluxy = np.zeros((imax, jmax+1), dtype=np.float64)
    gravity = float( 9.8 )
    
    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 )
    
    def flux(Qp, Qm, depp, depm, zbp, zbm) : 
        r = f(Qp, Qm)
        
        if ( (depm + zbm) < zbp) and (depp <= hbuf) : r = 0.0
        if ( (depp + zbp) < zbm) and (depm <= hbuf) : r = 0.0
            
        return r
        
    for i in numba.prange( 1, imax ):
        for j in numba.prange( jmax ):
            c, xm = (i,j), (i-1,j)
            fluxx[c] = flux(qx[c], qx[xm], dep[c], dep[xm], zb[c], zb[xm])
            
    for i in numba.prange( imax ):
        for j in numba.prange( 1, jmax ):
            c, ym = (i,j), (i,j-1)
            fluxy[c] = flux(qy[c], qy[ym], dep[c], dep[ym], zb[c], zb[ym])
            
# wall boundary 
#     fluxy[:,0] = 0.0 #fluxy[:,0]
#     fluxy[:,-1] = 0.0 #fluxy[:,0]
    
    n = 0
    for i in numba.prange(1, imax-1):
#         for j in numba.prange(jmax):
        for j in numba.prange(1, jmax-1):
            c, xp, yp = (i, j), (i+1, j), (i, j+1)
            depn[c] = dep[c] - dt*(fluxx[xp] - fluxx[c])/dx - dt*(fluxy[yp] - fluxy[c])/dy + dt*add[c]
            if depn[c] < hmin : 
                n += 1
                depn[c] = hmin
                
# upstream boundary                
    depn[0, :] = depn[1, :]
    depn[:, 0] = depn[:, 1]
    
# downstream boundary 
    depn[-1,:] = depn[-2,:]
    depn[:, -1] = depn[:,-2]
    
    return depn, n<jupyter_output><empty_output><jupyter_text>Momentum Equation$$    \begin{align}        \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}         - \nu_t h \left(\frac{\partial^2 u}{\partial x^2}+\frac{\partial^2 u}{\partial y^2} \right)= 0 \\        \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)        = 0    \end{align}$$<jupyter_code>@numba.jit(nopython=True, parallel=False)
def momentEq(dep, qx, qy, depn, zb, dt, dx, dy, cManning, hmin, hbuf, direction):
#direction = 1:x, 2:y
    gravity = float( 9.8 )
    
    q = qx if direction == 1 else qy
    u, v = qx/dep, qy/dep
    Vdir = q/dep  
        
    imax, jmax = len(q), len(q[0])
    
    qn = np.zeros_like(q, dtype=np.float64)
    fluxx = np.zeros((imax+1, jmax), dtype=np.float64)
    fluxy = np.zeros((imax, jmax+1), dtype=np.float64)
    
    f = lambda vp,vm,qp,qm : vm*qm if vp >= 0.0 and vm >= 0.0 else \
        (vp*qp if vp <= 0.0 and vm <= 0.0 else (0.5*vp+0.5*vm)*(0.5*qp+0.5*qm) )
        
    def flux1(vp, vm, qp, qm, depp, depm, zbp, zbm) : 
        if ( (depm + zbm) < zbp) and (depp <= hbuf) : r = 0.0
        if ( (depp + zbp) < zbm) and (depm <= hbuf) : r = 0.0
            
        r = f(vp,vm,qp,qm)
        return r
    
    for i in numba.prange( 1, imax ):
        for j in numba.prange( jmax ):
            c, xm = (i,j), (i-1,j)
            fluxx[c] = flux1( u[c], u[xm], q[c], q[xm], dep[c], dep[xm], zb[c], zb[xm] )
            
# boundary : not use
    fluxx[-1,:] = -9999.0
    fluxx[0,:]  = -9999.0

    for i in numba.prange( imax ):
        for j in numba.prange( 1, jmax ):
            c, ym = (i,j), (i,j-1)
            fluxy[c] = flux1( v[c], v[ym], q[c], q[ym], dep[c], dep[ym], zb[c], zb[ym] )
            
# wall boundary 
#     fluxy[:,0]  = -fluxy[:,1]
#     fluxy[:,-1] = -fluxy[:,-2]
    fluxy[:,0]  = -9999.0
    fluxy[:,-1] = -9999.0
    
    for i in numba.prange(1, imax-1):
#         for j in numba.prange(jmax):    
        for j in numba.prange(1, jmax-1):    
            c = (i, j)
            if depn[c] <= hbuf :
                qn[c] = 0.0
            else:
            # pressure & gravity term
                if direction == 2 and ((j == 0) or (j == jmax-1)) : 
                    if j == 0 :
                        c, yp = (i, j), (i, j+1)
                        Hc, Hp = depn[c] + zb[c], depn[yp] + zb[yp]
                        
                        if Hc < zb[yp] and depn[yp] <= hbuf :
                            dHdx = 0.0
                        else :
                            dHdx = ( Hp - Hc )/dy
                    elif j == jmax-1 :
                        c, ym = (i, j), (i, j-1)
                        Hc, Hm = depn[c] + zb[c], depn[ym] + zb[ym]
                        
                        if Hc < zb[ym] and depn[ym] <= hbuf :
                            dHdx = 0.0
                        else :
                            dHdx = ( Hc - Hm )/dy
                            
#                 elif direction == 1 and i == imax-1 :
#                     pass
                else :
                    c  = (i, j)
                    if direction == 1 :
                        xp = (i+1, j) #if i == imax-1 else (i+1, j)
                        xm = (i-1, j)
                        dp, dm, delta = xp, xm, dx
                    else :
                        yp = (i, j+1)
                        ym = (i, j-1)
                        dp, dm, delta = yp, ym, dy
                        
                    Vc, Vp, Vm = q[c]/dep[c], q[dp]/dep[dp], q[dm]/dep[dm]
                    Hc, Hp, Hm = depn[c]+zb[c], depn[dp]+zb[dp], depn[dm]+zb[dm]
                        
                    if Hc < zb[dp] and depn[dp] <= hbuf :
                        if Hc < zb[dm] and depn[dm] <= hbuf :
                            dHdx = 0.0
                        else:
                            dHdx = (Hc-Hm)/delta
                    elif Hc < zb[dm] and depn[dm] <= hbuf :
                        dHdx = (Hp-Hc)/delta
                    else :
                        if Vc > 0.0 and Vp > 0.0 and Vm > 0.0: 
                            Cr1, Cr2 = 0.5*(abs(Vc)+abs(Vp))*dt/delta, 0.5*(abs(Vc)+abs(Vm))*dt/delta
                            dHdx1, dHdx2 = (Hp-Hc)/delta, (Hc-Hm)/delta
                        elif Vc < 0.0 and Vp < 0.0 and Vm < 0.0:
                            Cr1, Cr2 = 0.5*(abs(Vc)+abs(Vm))*dt/delta, 0.5*(abs(Vc)+abs(Vp))*dt/delta
                            dHdx1, dHdx2 = (Hc-Hm)/delta, (Hp-Hc)/delta
                        else:
                            Cr1 = Cr2 = 0.5*(abs(0.5*(Vc+Vp))+abs(0.5*(Vc+Vm)))*dt/delta
                            dHdx1 = dHdx2 = (0.5*(Hc+Hp) - 0.5*(Hc+Hm)) / delta
            
                        w1, w2 = 1-Cr1**0.5, Cr2**0.5
                        dHdx = w1 * dHdx1 + w2 * dHdx2   
                
# viscous sublayer
                Cf = gravity*cManning**2.0/dep[c]**(1.0/3.0) 
                Vnorm = np.sqrt(u[c]**2.0+v[c]**2.0) 
                Vis = Cf * Vnorm * u[c] if direction == 1 else  Cf * Vnorm * v[c]
            
# turbulence
                nut = 0.4/6.0*dep[c]*np.sqrt(Cf)*np.abs(Vnorm)
                c  = (i, j)
                xp = (i+1, j) 
                xm = (i-1, j)
                yp = (i, j+1)
                ym = (i, j-1)
    
# side boundary : non-slip condition
#                 if i == imax-1: 
#                     turb = 0.0
                if j == 0 : 
                    if direction == 1:
                        turb = nut * ( Vdir[xp] - 2.0*Vdir[c] + Vdir[xm] )/ dx**2 \
                             + nut * ( Vdir[yp] - 3.0*Vdir[c] )/ dy**2
                    elif direction == 2:
                        turb = nut * ( Vdir[xp] - 2.0*Vdir[c] + Vdir[xm] )/ dx**2 \
                             + nut * ( Vdir[yp] - 3.0*Vdir[c] )/ dy**2
                elif j == jmax-1 : 
                    if direction == 1:
                        turb = nut * ( Vdir[xp] - 2.0*Vdir[c] + Vdir[xm] )/ dx**2 \
                             + nut * (  - 3.0*Vdir[c] + Vdir[ym] )/ dy**2
                    elif direction == 2:
                        turb = nut * ( Vdir[xp] - 2.0*Vdir[c] + Vdir[xm] )/ dx**2 \
                             + nut * (  - 3.0*Vdir[c] + Vdir[ym] )/ dy**2
                else :
                    turb = nut * ( Vdir[xp] - 2.0*Vdir[c] + Vdir[xm] )/ dx**2 \
                         + nut * ( Vdir[yp] - 2.0*Vdir[c] + Vdir[ym] )/ dy**2
                        
                c, xp, yp = (i,j), (i+1,j), (i,j+1)
                sourcet = Vis - dep[c] * turb
                
                qn[c] = q[c] - dt * ( fluxx[xp] - fluxx[c] ) / dx \
                             - dt * ( fluxy[yp] - fluxy[c] ) / dy \
                             - dt * gravity * depn[c] * dHdx \
                             - dt * sourcet
                            
# upstream boundary
#     if direction == 2 : 
#         qn[0,:] = 0.0
#     else :
#         updep = depn.copy()
#         for j in range(jmax):
#             updep[0,j] = 0.0 if updep[0,j] <= hmin else updep[0,j]
            
#         print(np.sum( updep[0,:]**(5/3) ))
#         alpha = Qup / dy / np.sum( updep[0,:]**(5/3) )
#         qn[0,:] = qup
#     alpha * updep[0,:]**(5/3)
        
# downstream boundary
    qn[-1,:] = qn[-2,:]
    qn[0,:] = qn[1,:]
    qn[:, -1] = qn[:, -2]
    qn[:,  0] = qn[:, 1]
        
    return qn
@numba.jit(nopython=True, parallel=False)
def simulation(dep, qx, qy, zb, dt, dx, dy, r, cManning, hmin, hbuf):
    depn, count = conEq(dep, qx, qy, zb, dt, dx, dy, r, hmin, hbuf)
    qxn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, cManning, hmin, hbuf, 1)
    qyn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, cManning, hmin, hbuf, 2)
#     CFL = ((np.abs(qxn/depn) + np.sqrt(9.8*depn))/dx + ( np.abs(qyn/depn) + np.sqrt(9.8*depn) )/dy)*dt
    CFL = ((np.abs(qxn/depn))/dx + ( np.abs(qyn/depn))/dy)*dt
    CFLmax = np.max(CFL)

    return depn, qxn, qyn, CFLmax, count<jupyter_output><empty_output><jupyter_text>read zb data<jupyter_code>ds = xr.open_dataset('zb.nc')
zb = np.copy(ds['elevation'].values)
zb = np.where( np.isnan(zb),999,zb)<jupyter_output><empty_output><jupyter_text>Main routine<jupyter_code>%%time

hmin = float(10.0**(-5)) 
hbuf = float(10.0**(-2)) 
dtout= float(900.0)
trunup = float(0.0 * 3600.0)
tmax = float(10.01 * 3600.0)
# tmax = float(1.0)

nout = 0
CFL = float(0.05)
cManning = 0.04

dx, dy = float(10), float(10)
dtini = float(0.1)

nxmax, nymax = ds.dims['x'], ds.dims['y']
qx = np.zeros((nxmax,nymax), dtype=np.float64)
qy = np.zeros_like(qx, dtype=np.float64)
dep = np.full_like(qx, hmin, dtype=np.float64)
r = np.zeros_like(qx, dtype=np.float64)
# zb[:,:] = ds['elevation'].values[:,:]

r[206:211,81] = 0.2

t = trunup
dt = dtini


while tmax >= t :
    t += dt
    dep, qx, qy, CFLmax, count = simulation(dep, qx, qy, zb, dt, dx, dy, r, cManning, hmin, hbuf)
        
# update dt
    dt = np.round( dt * CFL/CFLmax, 5)
    
    if t >= nout*dtout :
        print(t, dt, CFLmax, count)
        
        dss = xr.Dataset({'depth': (['x','y'], dep), 'u': (['x','y'], qx/dep), 'v': (['x','y'], qy/dep)
                          , 'elevation': (['x','y'], zb) }
                          , coords={'xc': (('x', 'y'), ds['xc']), 'yc': (('x', 'y'), ds['yc'])}
                          , attrs={'total_second' : round(t, 2)} )
        
        dss.to_netcdf('out' + str(nout).zfill(8) + '.nc')
        dss.close()
        nout += 1 
    
print(t, dt, CFLmax, count)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18432791852537664, 'fraction_numerical': 0.03206800774345594, 'mean_word_length': 1.3730776912322749, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '461580', 'license': 'MIT', 'n_tokens_mistral': 5008, 'n_tokens_neox': 4758, 'n_words': 1034, 'path': 'D2Dmodel-Copy1.ipynb', 'repo_name': 'computational-sediment-hyd/ChikumaRivOnOct2019-FloodWaterD2Dmodel', 'size': '11881'}