import module

In [1]:
import numpy as np
import numba
import xarray as xr

Continuous Equation

$$ \begin{align} \frac{\partial h}{\partial t}+\frac{\partial q_x}{\partial x} +\frac{\partial q_y}{\partial y} = 0 \end{align} $$
In [2]:
@numba.jit(nopython=True, parallel=False)
def conEq(dep, qx, qy, dt, dx, dy, hmin):
    imax, jmax = len(dep), len(dep[0])
    depn = np.empty_like(dep)
    fluxx = np.zeros((imax+1, jmax))
    fluxy = np.zeros((imax, jmax+1))
    
#     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 )
    def flux( Qp, Qm) : 
        if Qp > 0.0 and Qm > 0.0 :
            r = Qm
        elif Qp < 0.0 and Qm < 0.0 :
            r = Qp
        else :
# treatment of dry-bed
            if Qp == 0.0 and Qm == 0.0 :
                r = 0.0
            elif Qp == 0.0 :
                r = Qm if Qm > 0.0 else 0.0
            elif Qm == 0.0 :
                r = Qp if Qp < 0.0 else 0.0
            else :
                r = 0.5*Qp+0.5*Qm 
                
        return r
        
    
    for i in numba.prange( imax ):
        for j in numba.prange( jmax ):
            c, xm, ym = (i,j), (i-1,j), (i,j-1)
            fluxx[c] = flux(qx[c], qx[xm])
            fluxy[c] = flux(qy[c], qy[ym])
            
# continuous boundary 
    fluxx[-1,:] = fluxx[0,:]
    fluxy[:,-1] = fluxy[:,0]
    
    for i in numba.prange(imax):
        for j in numba.prange(jmax):
            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
            if depn[c] < hmin*2 : depn[c] = hmin

    return depn

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} $$
In [3]:
@numba.jit(nopython=True, parallel=False)
def momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, direction):
    #direction = 1:x, 2:y
    gravity = 9.8
    cManing = 0.02
    ib = 1.0/670.0

    q = qx if direction == 1 else qy
    u, v = qx/dep, qy/dep
        
    imax, jmax = len(q), len(q[0])
    
    qn = np.empty_like(q)
    fluxx = np.zeros((imax+1, jmax))
    fluxy = np.zeros((imax, jmax+1))

    flux = 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) )
    
    for i in numba.prange( imax ):
        for j in numba.prange( jmax ):
            c, xm, ym = (i,j), (i-1,j), (i,j-1)
            fluxx[c] = flux( u[c], u[xm], q[c], q[xm] )
            fluxy[c] = flux( v[c], v[ym], q[c], q[ym] )
            
# continuous boundary 
    fluxx[-1,:] = fluxx[0,:]
    fluxy[:,-1] = fluxy[:,0]
            
    for i in numba.prange(imax):
        for j in numba.prange(jmax):    
            c = (i, j)
            if depn[c] <= hmin*2 :
                qn[c] = 0.0
            else:
        # pressure & gravity term
                c, xm, ym = (i,j), (i-1, j), (i, j-1)
                xp = (0, j) if i == imax - 1 else (i+1, j  )
                yp = (i, 0) if j == jmax - 1 else (i  , j+1)
        
                if direction == 1 :
                    dp, dm, delta = xp, xm, dx 
                else :
                    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 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-ib, (Hc-Hm)/delta-ib
                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-ib, (Hp-Hc)/delta-ib          
                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 - ib
    
                w1, w2 = 1-Cr1**0.5, Cr2**0.5
                dHdx = w1 * dHdx1 + w2 * dHdx2   

#viscous sublayer
                Cf = gravity*cManing**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
#                kenergy = 2.07*Cf*Vnorm**2
                nut = 0.4/6.0*dep[c]*np.sqrt(Cf)*np.abs(Vnorm)
    
                turb = nut * ( u[xp] - 2.0*u[c] + u[xm] )/ dx**2 \
                     + nut * ( u[yp] - 2.0*u[c] + u[ym] )/ dy**2

            
                c, xp, yp = (i,j), (i+1,j), (i,j+1)
                qn[c] = q[c] - dt * ( fluxx[xp] - fluxx[c] ) / dx \
                             - dt * ( fluxy[yp] - fluxy[c] ) / dy \
                             - dt * gravity * depn[c] * dHdx \
                             - dt * Vis \
                             + dt * dep[c] * turb

    return qn
In [4]:
@numba.jit(nopython=True, parallel=False)
def simulation(dep, qx, qy, zb, dt, dx, dy, hmin):
    depn = conEq(dep, qx, qy, dt, dx, dy, hmin)
    qxn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, 1)
    qyn = momentEq(dep, qx, qy, depn, zb, dt, dx, dy, hmin, 2)
    CFL = ((np.abs(qxn/depn))/dx + (np.abs(qyn/depn))/dy)*dt
    CFLmin = np.max(CFL)
    return depn, qxn, qyn, CFLmin     

Main routine

read elevation data

In [5]:
ds = xr.open_dataset('dz.nc')
dza = ds['dz']
dz = dza.values
ds.close()

main

In [8]:
%%time

hmin = 10.0**(-3)
nxmax = len(dz)
nymax = len(dz[0])
dx = 5.0
dy = 5.0
dt = 0.1
dtout= 200.0
tmax = 3600.0

qx = np.zeros((nxmax,nymax))
qy = np.zeros_like(qx)
zb = dz #np.zeros_like(qx)

# Initial Condition
dep = np.full_like(zb, 0.4)

def timeGenerater():
    it, itout = 0, 1
    isTiter, isOut = True, False
    yield it*dt, isTiter, isOut

    while True:
        it += 1
        if it*dt >= itout*dtout:
            itout += 1
            isOut = True
        else:
            isOut = False

        if it*dt >= tmax: isTiter = False

        yield it*dt, isTiter, isOut
        
tgen = timeGenerater() 
t, isTiter, isOut = next(tgen)
tarr = [t]
deparr = np.array([dep])
uarr = np.array([qx/dep])
varr = np.array([qy/dep])

nout= 1
while isTiter:
    dep, qx, qy, CFLmin = simulation(dep, qx, qy, zb, dt, dx, dy, hmin)
    t, isTiter, isOut = next(tgen)
    if isOut:
        tarr.append(round(t,2))
        deparr = np.append(deparr, np.array([dep]), axis=0)
        uarr = np.append(uarr, np.array([qx/dep]), axis=0)
        varr = np.append(varr, np.array([qy/dep]), axis=0)
        print(t,CFLmin)
        nout += 1
        
dss = xr.Dataset({'depth': (['t','x','y'], deparr), 'u': (['t','x','y'], uarr), 'v': (['t','x','y'], varr)}, coords={'xc': (('x', 'y'), dza['xc'])
                                             , 'yc': (('x', 'y'), dza['yc'])
                                             , 't': tarr})

dss.to_netcdf('output.nc')
200.0 0.12994759757305765
400.0 0.05324376553275483
600.0 0.13465034344958737
800.0 0.08917062844146038
1000.0 0.06319218033460254
1200.0 0.11995618722432487
1400.0 0.07224942603327707
1600.0 0.12386222201873731
1800.0 0.06899316890484393
2000.0 0.170605601280248
2200.0 0.06517755577401307
2400.0 0.1353129352137529
2600.0 0.0615231283218484
2800.0 0.07240785399189434
3000.0 0.07694049275220537
3200.0 0.06135978323502305
3400.0 0.08224366525350822
3600.0 0.061249249226995206
Wall time: 21.4 s

check total volume

In [9]:
np.sum( deparr[-1] ), np.sum( deparr[0] )
Out[9]:
(1524.5068734878676, 1527.6000000000004)