import numpy as np
import numba
import xarray as xr
@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
@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
@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
ds = xr.open_dataset('dz.nc')
dza = ds['dz']
dz = dza.values
ds.close()
%%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
np.sum( deparr[-1] ), np.sum( deparr[0] )
(1524.5068734878676, 1527.6000000000004)