@numba.jit(nopython=True, parallel=False)
def flowmodel(A, Q, Adown, Qup, zb, B, dx, dt, g, manning):
imax = len(A)
Anew, Qnew = np.zeros(imax), np.zeros(imax)
# continuous equation
for i in numba.prange(1, imax-1) : Anew[i] = A[i] - dt * ( Q[i] - Q[i-1] ) / dx
Anew[0], Anew[-1] = Anew[1], Adown
# moumentum equation
for i in numba.prange(1,imax-1):
ip, ic, im = (i+1, i, i-1)
Cr1 = 0.5*( Q[ic]/A[ic] + Q[ip]/A[ip] )*dt/dx
Cr2 = 0.5*( Q[ic]/A[ic] + Q[im]/A[im] )*dt/dx
dHdx1 = ( Anew[ip]/B[ip] + zb[ip] - Anew[ic]/B[ic] - zb[ic] ) / dx
dHdx2 = ( Anew[ic]/B[ic] + zb[ic] - Anew[im]/B[im] - zb[im] ) / dx
dHdx = (1.0 - Cr1) * dHdx1 + Cr2 * dHdx2
Qnew[ic] = Q[ic] - dt * ( Q[ic]**2/A[ic] - Q[im]**2/A[im] ) / dx \
- dt * g * Anew[ic] * dHdx \
- dt * g * A[ic] * manning**2 * Q[ic]**2 / B[ic]**2 / ( A[ic]/B[ic] )**(10.0/3.0)
# Qnew[0], Qnew[-1] = Qup, Qnew[-2]
Qnew[0] = Qup
# check downstream boundary Q
i = -1
ic, im = (i, i-1)
dHdx = ( Anew[ic]/B[ic] + zb[ic] - Anew[im]/B[im] - zb[im] ) / dx
Qnew[ic] = Q[ic] - dt * ( Q[ic]**2/A[ic] - Q[im]**2/A[im] ) / dx \
- dt * g * Anew[ic] * dHdx \
- dt * g * A[ic] * manning**2 * Q[ic]**2 / B[ic]**2 / ( A[ic]/B[ic] )**(10.0/3.0)
return Anew, Qnew