def umultiply(pow, N, u):
""" Compute the sum of all possible combinations yielding the power 'pow' of signal 'u' with a total order 'N'
i.e. (u^pow)^<N>
"""
v = 0
if pow>2:
for i in range(0, N+1):
v += ny.complexAmplitudeProduct(umultiply(2, i, u), umultiply(pow-2, N-i, u), 2)
else:
for i in range(0, N+1):
v += ny.complexAmplitudeProduct(u[i], u[N-i], 2)
return v
def surfaceForcingDer(self, order, der):
"""Calculate forcing F in equation
u_t^<order>(der)-( Av u_z^<order> )^(der+1) = F^<order>(der) at surface z=0
"""
jmax = self.input.v('grid', 'maxIndex', 'x')
fmax = self.input.v('grid', 'maxIndex', 'f')
forcing = np.zeros((jmax + 1, 1, fmax + 1), dtype=complex)
# forcing from barotropic pressure gradient
if der == 0:
forcing += self.input.v('G') * self.input.d('zeta' + str(order),
range(0, jmax + 1),
[0],
range(0, fmax + 1),
dim='x')
# forcing by advection; only include this in the reconstruction if the equation at 'order' contained advection
if 'adv' in self.input.getKeysOf('u' + str(order)):
for order1 in range(0, order):
order2 = order - order1 - 1
for i in range(0, der + 1):
u = self.surf_u_der[:, :, :, order1, i]
ux = ny.derivative(
self.surf_u_der[:, :, :, order2, der - i], 'x',
self.input.slice('grid'))
forcing += scipy.misc.comb(
der, i) * ny.complexAmplitudeProduct(u, ux, 2)
if i == 0:
wstr = 'w' + str(order1)
w = self.input.v(wstr, range(0, jmax + 1), [0],
range(0, fmax + 1))
else:
Bu = self.input.v('B', range(
0, jmax + 1), [0], range(
0, fmax + 1)) * self.surf_u_der[:, :, :,
order1, i - 1]
w = ny.derivative(
Bu, 'x', self.input.slice('grid')) / self.input.v(
'B', range(0, jmax + 1), [0], range(
0, fmax + 1))
uz = self.surf_u_der[:, :, :, order2, der - i + 1]
forcing += scipy.misc.comb(
der, i) * ny.complexAmplitudeProduct(w, uz, 2)
# TODO
#if 'baroc' in self.submodulesToRun:
#if 'densitydrift' in self.submodulesToRun:
return forcing
def erosion(ws,
tau_order,
data,
method='Chernetsky',
submodule=None,
friction='Roughness'):
jmax = data.v('grid', 'maxIndex', 'x')
kmax = data.v('grid', 'maxIndex', 'z')
fmax = data.v('grid', 'maxIndex', 'f')
rho0 = data.v('RHO0')
taub_abs = shearstress(tau_order,
data,
submodule=submodule,
friction=friction)
## 2. erosion
finf = data.v('finf')
if method == 'Partheniades':
hatE = finf * taub_abs
else:
rhos = data.v('RHOS')
gred = data.v('G') * (rhos - rho0) / rho0
ds = data.v('DS')
hatE = finf * rhos / (gred * ds * rho0) * ny.complexAmplitudeProduct(
ws[:, [kmax], :], taub_abs, 2)
return hatE[:, :, :fmax + 1]
def updateSurfaceData(self):
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
OMEGA = self.input.v('OMEGA')
# 1. update velocity and its derivative of previous order
ustr = 'u' + str(self.currentOrder - 1)
self.surf_u_der[:, :, :, self.currentOrder - 1,
0] = self.input.v(ustr, range(jmax + 1), [0],
range(fmax + 1))
self.surf_u_der[:, :, :, self.currentOrder - 1,
1] = self.input.d(ustr,
range(jmax + 1), [0],
range(fmax + 1),
dim='z')
# 2. calculate other terms
for order in range(0, self.currentOrder):
der = self.currentOrder - order
# 2a. update surface stress term
forcing = self.surfaceForcingDer(order, der - 1)
D = (np.arange(0, fmax + 1) * 1j * OMEGA).reshape(
(1, 1, fmax + 1)) * np.ones((jmax + 1, 1, 1))
self.surf_stress[:, :, :, order,
der - 1] = D * self.surf_u_der[:, :, :, order,
der - 1] + forcing
# 2b. update surface velocity derivative
sum_nu_u = np.zeros((jmax + 1, 1, fmax + 1), dtype=complex)
for i in range(1, der + 1):
if i == 1:
Avder = self.input.d('Av',
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='z') # first der of Av
elif i == 2:
Av = self.input.d('Av',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='zz')
Avder = Av[:, [0], :] # 2nd der of Av
else:
Av = ny.derivative(Av, 'z', self.input.slice('grid'))
Avder = Av[:,
[0], :] # every step: take a higher der of Av
sum_nu_u += scipy.misc.comb(
der, i) * ny.complexAmplitudeProduct(
Avder, self.surf_u_der[:, :, :, order, der + 1 - i], 2)
Av = self.input.v('Av', range(0, jmax + 1), [0],
range(0, fmax + 1))
self.surf_u_der[:, :, :, order, der + 1] = -self.Avinv_multiply(
Av, sum_nu_u - self.surf_stress[:, :, :, order, der - 1])
return
def shearstressGS(tau_order, data, submodule=None, friction='Roughness'): # Shear stress following the formulation of Schramkowski; only subtidal sf and subtidal second order friction
jmax = data.v('grid', 'maxIndex', 'x')
kmax = data.v('grid', 'maxIndex', 'z')
fmax = data.v('grid', 'maxIndex', 'f')
rho0 = data.v('RHO0')
sf = data.v(friction, range(0, jmax+1), 0, 0)
if submodule is None:
submodule = (None, )*(tau_order+1)
## 1. bed shear stress
# the bed shear stress is extended over fmax+1 frequency components to prevent inaccuracies in truncation
ulist = []
for i in range(0, tau_order+1):
if submodule[i] is None:
u = data.v('u'+str(i), range(0, jmax+1), [kmax], range(0, fmax+1))
else:
u = data.v('u'+str(i), submodule[i], range(0, jmax+1), [kmax], range(0, fmax+1))
if u is None:
u = np.zeros((jmax+1, 1, fmax+1), dtype=complex)
ulist.append(u)
taub_abs = np.zeros((jmax+1, 1, fmax+1), dtype=complex)
if tau_order == 0:
# uabs0 = ny.absoluteU(ulist[0][:, 0, 1]+10**-6, 0)
# uabs2 = ny.absoluteU(ulist[0][:, 0, 1]+10**-6, 2)+np.conj(ny.absoluteU(ulist[0][:, 0, 1]+10**-6, -2))
uabs0 = ny.absoluteU(ulist[0][:, 0, 1], 0)
uabs2 = ny.absoluteU(ulist[0][:, 0, 1], 2) + np.conj(ny.absoluteU(ulist[0][:, 0, 1], -2))
taub_abs[:, 0, 0] = rho0*sf*uabs0
taub_abs[:, 0, 2] = rho0*sf*uabs2
elif tau_order ==1:
signu = np.zeros((jmax+1, 1, np.maximum(fmax+1, 4)), dtype=complex)
# signu[:, 0, 1] = ny.signU(ulist[0][:, 0, 1]+10**-6, 1) + np.conj(ny.signU(ulist[0][:, 0, 1]+10**-6, -1))
# signu[:, 0, 3] = ny.signU(ulist[0][:, 0, 1]+10**-6, 3) + np.conj(ny.signU(ulist[0][:, 0, 1]+10**-6, -3))
signu[:, 0, 1] = ny.signU(ulist[0][:, 0, 1], 1) + np.conj(ny.signU(ulist[0][:, 0, 1], -1))
signu[:, 0, 3] = ny.signU(ulist[0][:, 0, 1], 3) + np.conj(ny.signU(ulist[0][:, 0, 1], -3))
if fmax+1 < 4:
ulist[1] = np.concatenate((ulist[1], np.zeros((jmax+1, 1, 4-fmax-1))), 2)
taub_abs = rho0*sf.reshape((jmax+1, 1, 1))*ny.complexAmplitudeProduct(ulist[1], signu, 2)
elif tau_order ==2:
utid = ny.invfft2(ulist[0], 2, 90)
ucomb = ny.invfft2(ulist[0]+ulist[1], 2, 90)
uabs_tid = np.mean(np.abs(utid), axis=2)
uabs_tot = np.mean(np.abs(ucomb), axis=2)
uabs_eps = (uabs_tot - uabs_tid).reshape((jmax+1))
taub_abs[:, 0, 0] = rho0*sf*uabs_eps
return taub_abs[:, :, :fmax+1]
def shearstressCheb(tau_order, data, submodule=None, friction='Roughness'): # Shear stress using Chebyshev polynomials
jmax = data.v('grid', 'maxIndex', 'x')
kmax = data.v('grid', 'maxIndex', 'z')
fmax = data.v('grid', 'maxIndex', 'f')
fmax = np.maximum(fmax, 3) ## NB. take fmax 3 at minimum: erosion uses |u0|*u1 which, with standard assumptions, yields an M2 signal for |u0| containing an M2 and M6
if submodule is None:
submodule = (None, )*(tau_order+1)
## 1. bed shear stress
# the bed shear stress is extended over fmax+1 frequency components to prevent inaccuracies in truncation
# Method 1: using Av
taub1 = []
for i in range(0, tau_order+1):
taub1.append(0)
for j in range(0, i+1):
# friction
if j == 0:
Av = data.v('Av', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
elif submodule[i] is None or submodule[i]== 'mixing':
Av = data.v('Av'+str(i), range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
else:
Av = np.zeros((jmax+1, kmax+1, fmax+1))
# velocity shear
q = i-j
if submodule[q] is None:
uz = data.d('u'+str(q), range(0, jmax+1), [kmax], range(0, fmax+1), dim='z')
else:
uz = data.d('u'+str(q), submodule[q], range(0, jmax+1), [kmax], range(0, fmax+1), dim='z')
if uz is None:
uz = np.zeros((jmax+1, 1, fmax+1), dtype=complex)
# extend and multiply
Av = np.concatenate((Av, np.zeros((jmax+1, kmax+1, fmax+1))), 2)
uz = np.concatenate((uz, np.zeros((jmax+1, 1, fmax+1))),2)
taub1[i] += ny.complexAmplitudeProduct(Av[:, [kmax], :], uz, 2)
# Method 2: using sf (equivalent to using Av)
taub = []
for i in range(0, tau_order+1):
taub.append(0)
for j in range(0, i+1):
# friction
if j == 0:
sf = data.v(friction, range(0, jmax+1), [0], range(0, fmax+1))
elif submodule[i] is None or submodule[i]== 'mixing':
sf = data.v(friction+str(i), range(0, jmax+1), [0], range(0, fmax+1))
else:
sf = np.zeros((jmax+1, 1, fmax+1))
# velocity
q = i-j
if submodule[q] is None:
u = data.v('u'+str(q), range(0, jmax+1), [kmax], range(0, fmax+1))
else:
u = data.v('u'+str(q), submodule[q], range(0, jmax+1), [kmax], range(0, fmax+1))
if u is None:
u = np.zeros((jmax+1, 1, fmax+1), dtype=complex)
# extend vectors and multiply
sf = np.concatenate((sf, np.zeros((jmax+1, 1, fmax+1))), 2)
u = np.concatenate((u, np.zeros((jmax+1, 1, fmax+1))),2)
taub[i] += ny.complexAmplitudeProduct(sf, u, 2)
taub = taub1
for i, t in enumerate(taub):
dif = np.max(np.abs(t - taub1[i]))
if dif>1.e-15:
print 'test stress ' + str(np.max(np.abs(t - taub1[i])))
print 'order ' + str(tau_order)
# amplitude
tau_amp = (np.sum(np.abs(sum(taub)), axis=-1)+10**-3).reshape((jmax+1, 1, 1))
taub = [t/tau_amp for t in taub]
# absolute value
c = ny.polyApproximation(np.abs, 8) # chebyshev coefficients for abs
taub_abs = np.zeros(taub[0].shape, dtype=complex)
if tau_order==0:
taub_abs[:, :, 0] = c[0]
taub_abs += c[2]*umultiply(2, tau_order, taub)
taub_abs += c[4]*umultiply(4, tau_order, taub)
taub_abs += c[6]*umultiply(6, tau_order, taub)
taub_abs += c[8]*umultiply(8, tau_order, taub)
rho0 = data.v('RHO0')
taub_abs = taub_abs*tau_amp*rho0
return taub_abs[:, :, :fmax+1]
def run(self):
self.logger.info('Running module SalinityFirst')
# Init
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
SIGMASAL = self.input.v('SIGMASAL')
OMEGA = self.input.v('OMEGA')
submodulesToRun = self.input.v('submodules')
sx0 = self.input.d('s0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
u0 = self.input.v('u0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
s1var = self.input.v('s1var', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
zeta0 = self.input.v('zeta0', range(0, jmax + 1), [0],
range(0, fmax + 1))
H = self.input.v('H', range(0, jmax + 1))
B = self.input.v('B', range(0, jmax + 1))
AKh = self.input.v('Kh', range(0, jmax + 1)) * B * H
################################################################################################################
# Second-order salinity variation as function of first-order salinity closure
################################################################################################################
# build, save and solve the velocity matrices in every water column
# LHS
Kv = self.input.v('salinityMatrix')
hasMatrix = True
if Kv is None:
Kv = self.input.v('Av', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1)) / SIGMASAL
hasMatrix = False
# RHS
# Allow for three submodules: advection, diffusion and nostress
nRHS = len(self.input.getKeysOf('u1')) + 3
f_index = -1
f_names = []
F = np.zeros([jmax + 1, kmax + 1, fmax + 1, nRHS], dtype=complex)
Fsurf = np.zeros([jmax + 1, 1, fmax + 1, nRHS], dtype=complex)
Fbed = np.zeros([jmax + 1, 1, fmax + 1, nRHS], dtype=complex)
if 'advection' in submodulesToRun:
# advection by u0*sx1
sx1var = self.input.d('s1var',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
sz1var = self.input.d('s1var',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
w0 = self.input.v('w0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
f_index += 1
f_names.append(['advection', 'u0-s1'])
F[:, :, :, f_index] = -ny.complexAmplitudeProduct(
u0, sx1var, 2) - ny.complexAmplitudeProduct(w0, sz1var, 2)
del sx1var, sz1var, w0
# advection by u1*sx0
submods = self.input.getKeysOf('u1')
for mod in submods:
u1 = self.input.v('u1', mod, range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
f_index += 1
f_names.append(['advection', 'u1_' + mod + '-s0'])
F[:, :, :, f_index] = -ny.complexAmplitudeProduct(u1, sx0, 2)
del u1
if 'diffusion' in submodulesToRun:
f_index += 1
f_names.append(['diffusion', 's0'])
F[:, :, 0, f_index] = (ny.derivative(AKh * sx0[:, 0, 0], 'x',
self.input.slice('grid')) /
(B * H)).reshape((jmax + 1, 1)) * np.ones(
(1, kmax + 1))
if 'nostress' in submodulesToRun:
D = (np.arange(0, fmax + 1) * 1j * OMEGA).reshape(
(1, 1, fmax + 1)) * np.ones((jmax + 1, 1, 1))
Kvsz1z = D * s1var[:, [0], :] + ny.complexAmplitudeProduct(
u0[:, [0], :], sx0[:, [0], :], 2)
f_index += 1
f_names.append(['nostress', 's1-zeta0'])
Fsurf[:, 0, :, f_index] = -ny.complexAmplitudeProduct(
Kvsz1z, zeta0, 2).reshape((jmax + 1, fmax + 1))
del D, Kvsz1z
sCoef, sForced, szCoef, szForced, salinityMatrix = svarFunction(
Kv, F, Fsurf, Fbed, self.input, hasMatrix=hasMatrix)
del Kv
################################################################################################################
# First-order salinity closure
################################################################################################################
## LHS terms
# First-order river discharge
Q = -self.input.v('Q1', range(0, jmax + 1))
# Diffusion coefficient
us = ny.complexAmplitudeProduct(u0, sCoef,
2)[:, :, 0, 0] # subtidal part of u*s
us = ny.integrate(us, 'z', kmax, 0,
self.input.slice('grid')).reshape(jmax + 1)
AK = np.real(AKh) - np.real(B * us)
del us
## RHS terms
nRHS_clo = nRHS + len(self.input.getKeysOf('u1')) + 3
f_index_clo = -1
f_names_clo = []
F = np.zeros([jmax + 1, nRHS_clo])
Fopen = np.zeros([1, nRHS_clo])
Fclosed = np.zeros([1, nRHS_clo])
if 'advection' in submodulesToRun:
# advection by u0*s2
us = ny.complexAmplitudeProduct(u0, sForced, 2)[:, :, [0], :]
us = ny.integrate(us, 'z', kmax, 0,
self.input.slice('grid')).reshape(
jmax + 1, nRHS)
for i in range(0, nRHS):
f_index_clo += 1
f_names_clo.append(['advection', 'u0-s2_' + f_names[i][0]])
F[:, f_index_clo] = -ny.derivative(np.real(B * us[:, i]), 'x',
self.input.slice('grid'))
# advection by u1*s1
submods = self.input.getKeysOf('u1')
for mod in submods:
u1 = self.input.v('u1', mod, range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
us = ny.complexAmplitudeProduct(u1, s1var, 2)[:, :, 0]
us = ny.integrate(us, 'z', kmax, 0,
self.input.slice('grid')).reshape(jmax + 1)
f_index_clo += 1
f_names_clo.append(['advection', 'u1_' + mod + '-s1'])
F[:, f_index_clo] = -ny.derivative(np.real(B * us), 'x',
self.input.slice('grid'))
del u1
# surface term
us = ny.complexAmplitudeProduct(u0[:, [0], :], s1var[:, [0], :], 2)
us = ny.complexAmplitudeProduct(us[:, [0], :], zeta0, 2)[:, 0, 0]
f_index_clo += 1
f_names_clo.append(['advection', 'surface'])
F[:, f_index_clo] = -ny.derivative(np.real(B * us), 'x',
self.input.slice('grid'))
del us
if 'diffusion' in submodulesToRun:
# Bed terms
Hx = self.input.d('H', range(0, jmax + 1), dim='x')
sx1var = self.input.d('s1var',
range(0, jmax + 1),
kmax,
0,
dim='x')
f_index_clo += 1
f_names_clo.append(['diffusion', 'bedslope'])
F[:, f_index_clo] = - ny.derivative(np.real(AKh*s1var[:, -1, 0]*Hx), 'x', self.input.slice('grid'))/H \
- np.real(AKh/H*sx1var*Hx)
# Surface term
f_index_clo += 1
f_names_clo.append(['diffusion', 'surface'])
F[:, f_index_clo] = np.real(zeta0[:, 0, 0]) * ny.derivative(
np.real(AKh * sx0[:, 0, 0]), 'x', self.input.slice('grid')) / H
del sx1var, Hx
## Solve equation
S1 = np.zeros((jmax + 1, 1, fmax + 1, nRHS_clo))
Sx1 = np.zeros((jmax + 1, 1, fmax + 1, nRHS_clo))
S1[:, 0, 0, :], Sx1[:, 0, 0, :] = sclosureFunction((Q, AK), F, Fopen,
Fclosed, self.input)
################################################################################################################
# Second-order salinity variation
################################################################################################################
s2 = np.zeros((jmax + 1, kmax + 1, fmax + 1, nRHS + 1), dtype=complex)
sz2 = np.zeros((jmax + 1, kmax + 1, fmax + 1, nRHS + 1), dtype=complex)
s2[:, :, :, :-1] = sForced
sz2[:, :, :, :-1] = szForced
f_names.append(['advection', 'u0-S1'])
s2[:, :, :, -1] = np.sum(ny.complexAmplitudeProduct(sCoef, Sx1, 2), 3)
sz2[:, :, :, -1] = np.sum(ny.complexAmplitudeProduct(szCoef, Sx1, 2),
3)
################################################################################################################
# Make final dictionary to return
################################################################################################################
d = {}
d['salinityMatrix'] = salinityMatrix
d['s1'] = {}
d['s2var'] = {}
for submod in submodulesToRun:
if submod in zip(*f_names_clo)[0]:
d['s1'][submod] = {}
if submod in zip(*f_names)[0]:
d['s2var'][submod] = {}
# if submod in zip(*f_names)[0]:
for i, mod in enumerate(f_names_clo):
nf = ny.functionTemplates.NumericalFunctionWrapper(
S1[:, :, :, i], self.input.slice('grid'))
nf.addDerivative(Sx1[:, :, :, i], 'x')
if len(mod) == 1:
d['s1'][mod[0]] = nf.function
if len(mod) == 2:
d['s1'][mod[0]][mod[1]] = nf.function
for i, mod in enumerate(f_names):
nf2 = ny.functionTemplates.NumericalFunctionWrapper(
s2[:, :, :, i], self.input.slice('grid'))
nf2.addDerivative(sz2[:, :, :, i], 'z')
if len(mod) == 1:
d['s2var'][mod[0]] = nf2.function
if len(mod) == 2:
d['s2var'][mod[0]][mod[1]] = nf2.function
return d
def component(self, ws, Kv, Kh):
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
z = ny.dimensionalAxis(self.input.slice('grid'), 'z')[:, :, 0]
depth = (self.input.v('grid', 'low', 'z', range(0, jmax+1)) - self.input.v('grid', 'high', 'z', range(0, jmax+1))).reshape((jmax+1, 1))
B = self.input.v('B', range(0, jmax+1))
u0 = self.input.v('u0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
w0 = self.input.v('w0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
zeta0 = self.input.v('zeta0', range(0, jmax+1), [0], range(0, fmax+1))
################################################################################################################
# Leading order
################################################################################################################
Chat = np.zeros((jmax+1, kmax+1, fmax+1), dtype=complex)
if ws[0,0] != 0:
k = depth*ws[:, [-1]]/Kv[:, [-1]]*(1-np.exp(-ws[:, [-1]]/Kv[:, [-1]]*depth))**-1.
else:
k = 1
Chat[:, :, 0] = k*np.exp(-ws[:, [-1]]/Kv[:, [-1]]*(z+depth)) # k is such that depth-av Chat = 1.
Chatx = ny.derivative(Chat, 'x', self.input.slice('grid'))
Chatz = -(ws[:, [-1]]/Kv[:, [-1]]).reshape((jmax+1, 1, 1))*Chat
################################################################################################################
# First order
################################################################################################################
F = np.zeros((jmax+1, kmax+1, 2*fmax+1, 2), dtype=complex)
Fsurf = np.zeros((jmax+1, 1, 2*fmax+1, 2), dtype=complex)
Fbed = np.zeros((jmax+1, 1, 2*fmax+1, 2), dtype=complex)
## forcing terms
# advection
F[:, :, fmax:, 0] = -ny.complexAmplitudeProduct(u0, Chatx, 2) - ny.complexAmplitudeProduct(w0, Chatz, 2)
F[:, :, fmax:, 1] = -ny.complexAmplitudeProduct(u0, Chat, 2)
## solve
Chat1, _ = pFunction(1, ws, Kv, F[:, :, fmax+1], Fsurf[:, :, fmax+1], Fbed[:, :, fmax+1], self.input)
Chat1 = ny.eliminateNegativeFourier(Chat1, 2)
################################################################################################################
# Closure
################################################################################################################
# transport
T = {}
T['adv'] = {}
T['adv']['tide'] = np.real(ny.integrate(ny.complexAmplitudeProduct(u0, Chat1[:, :, :, 0], 2), 'z', kmax, 0, self.input.slice('grid'))[:, 0, 0]*B)
for key in self.input.getKeysOf('u1'):
utemp = self.input.v('u1', key, range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
try:
T['adv'][key] += np.real(ny.integrate(ny.complexAmplitudeProduct(utemp, Chat, 2), 'z', kmax, 0, self.input.slice('grid'))[:, 0, 0]*B)
except:
T['adv'][key] = np.real(ny.integrate(ny.complexAmplitudeProduct(utemp, Chat, 2), 'z', kmax, 0, self.input.slice('grid'))[:, 0, 0]*B)
T['dif'] = - np.real(ny.integrate(ny.complexAmplitudeProduct(Kh, Chatx, 2), 'z', kmax, 0, self.input.slice('grid'))[:, 0, 0]*B)
T['noflux'] = np.real(ny.complexAmplitudeProduct(ny.complexAmplitudeProduct(u0[:, [0], :], Chat[:, [0], :], 2), zeta0, 2)[:, 0, 0]*B)
F = {}
F['adv'] = {}
F['adv']['tide'] = np.real(ny.integrate(ny.complexAmplitudeProduct(u0, Chat1[:, :, :, -1], 2), 'z', kmax, 0, self.input.slice('grid'))[:, 0, 0]*B)
F['dif'] = - np.real(Kh[:, 0]*depth[:, 0]*B)
H1 = np.real(depth[:, 0]*B)
return T, F, H1, Chat[:, :, 0]
def run(self):
self.logger.info('Running module SedDynamic - first order')
d = {}
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
ftot = 2*fmax+1
OMEGA = self.input.v('OMEGA')
self.submodulesToRun = ny.toList(self.input.v('submodules'))
method = self.input.v('erosion_formulation')
frictionpar = self.input.v('friction') # friction parameter used for the erosion, by default the total roughness
if frictionpar == None:
frictionpar = 'Roughness'
################################################################################################################
# Left hand side
################################################################################################################
Av = self.input.v('Av', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
Kv = self.input.v('Kv', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
# NB. If Kv is not provided on input, use the module DiffusivityUndamped to compute it. This is a fix for making this module easier to use.
if Kv is None:
from DiffusivityUndamped import DiffusivityUndamped
sr = self.input.v('sigma_rho')
if sr is None: # add Prandtl-Schmidt number if it does not exist
self.input.addData('sigma_rho', 1.)
md = DiffusivityUndamped(self.input)
self.input.merge(md.run())
Kv = self.input.v('Kv', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
d['Kv'] = Kv
ws = self.input.v('ws0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
################################################################################################################
# Forcing terms
################################################################################################################
if 'sedadv' in self.submodulesToRun and 'sedadv_ax' not in self.submodulesToRun:
self.submodulesToRun.append('sedadv_ax')
# determine number of submodules
nRHS = len(self.submodulesToRun)
if 'erosion' in self.submodulesToRun:
keysu1 = self.input.getKeysOf('u1')
try:
keysu1.remove('mixing') # flow due to 'mixing' should not be included in the erosion term
except:
pass
self.submodulesToRun.remove('erosion') #move to the end of the list
self.submodulesToRun.append('erosion')
nRHS = nRHS - 1 + len(self.input.getKeysOf('u1'))
F = np.zeros([jmax+1, kmax+1, ftot, nRHS], dtype=complex)
Fsurf = np.zeros([jmax+1, 1, ftot, nRHS], dtype=complex)
Fbed = np.zeros([jmax+1, 1, ftot, nRHS], dtype=complex)
c0 = self.input.v('hatc0')
cx0 = self.input.d('hatc0', dim='x')
cz0 = self.input.d('hatc0', dim='z')
# 1. Erosion
if 'erosion' in self.submodulesToRun:
# erosion due to first-order bed shear stress
# E = erosion(ws, Av, 1, self.input, method) # 24-04-2017 Obsolete
# Fbed[:, :, fmax:, self.submodulesToRun.index('erosion')] = -E # 24-04-2017 Obsolete
for submod in keysu1:
E = erosion(ws, 1, self.input, method, submodule=(None, submod), friction=frictionpar)
Fbed[:, :, fmax:, len(self.submodulesToRun)-1 + keysu1.index(submod)] = -E
# 2. Advection
if 'sedadv' in self.submodulesToRun:
u0 = self.input.v('u0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
w0 = self.input.v('w0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
eta = ny.complexAmplitudeProduct(u0, cx0, 2)+ny.complexAmplitudeProduct(w0, cz0, 2)
F[:, :, fmax:, self.submodulesToRun.index('sedadv')] = -eta
F[:, :, fmax:, self.submodulesToRun.index('sedadv_ax')] = -ny.complexAmplitudeProduct(u0, c0, 2)
# 3. First-order fall velocity
if 'fallvel' in self.submodulesToRun:
# surface and internal terms
ws1 = self.input.v('ws1', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
ksi = ny.complexAmplitudeProduct(ws1, c0, 2)
ksiz = ny.derivative(ksi, 'z', self.input.slice('grid'))
zeta0 = self.input.v('zeta0', range(0, jmax+1), 0, range(0, fmax+1))
F[:, :, fmax:, self.submodulesToRun.index('fallvel')] = ksiz
Fsurf[:, 0, fmax:, self.submodulesToRun.index('fallvel')] = -ny.complexAmplitudeProduct(ksiz[:,0,:], zeta0, 1)
# adjustment to erosion
E = erosion(ws1, 0, self.input, method, friction=frictionpar)
Fbed[:, :, fmax:, self.submodulesToRun.index('fallvel')] = -E
# 4. First-order eddy diffusivity
if 'mixing' in self.submodulesToRun:
# surface, bed and internal terms
Kv1 = self.input.v('Kv1', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
psi = ny.complexAmplitudeProduct(Kv1, cz0, 2)
psiz = ny.derivative(psi, 'z', self.input.slice('grid'))
F[:, :, fmax:, self.submodulesToRun.index('mixing')] = psiz
Fsurf[:, 0, fmax:, self.submodulesToRun.index('mixing')] = -psi[:, 0, :]
Fbed[:, 0, fmax:, self.submodulesToRun.index('mixing')] = -psi[:, -1, :]
# adjustment to erosion
E = erosion(ws, 1, self.input, method, submodule=(None, 'mixing'), friction=frictionpar)
Fbed[:, :, fmax:, self.submodulesToRun.index('mixing')] = -E
# 5. No-flux surface correction
if 'noflux' in self.submodulesToRun:
zeta0 = self.input.v('zeta0', range(0, jmax+1), [0], range(0, fmax+1))
D = np.zeros((jmax+1, 1, fmax+1, fmax+1), dtype=complex)
D[:, :, range(0, fmax+1), range(0, fmax+1)] = np.arange(0, fmax+1)*1j*OMEGA
Dc0 = ny.arraydot(D, c0[:, [0], :])
chi = ny.complexAmplitudeProduct(Dc0, zeta0, 2)
Fsurf[:, :, fmax:, self.submodulesToRun.index('noflux')] = -chi
################################################################################################################
# Solve equation
################################################################################################################
cmatrix = self.input.v('cMatrix')
if cmatrix is not None:
c, cMatrix = cFunction(None, cmatrix, F, Fsurf, Fbed, self.input, hasMatrix = True)
else:
c, cMatrix = cFunction(ws, Kv, F, Fsurf, Fbed, self.input, hasMatrix = False)
c = c.reshape((jmax+1, kmax+1, ftot, nRHS))
c = ny.eliminateNegativeFourier(c, 2)
################################################################################################################
# Prepare output
################################################################################################################
d['hatc1'] = {}
d['hatc1']['a'] = {}
d['hatc1']['ax'] = {}
for i, submod in enumerate(self.submodulesToRun):
if submod == 'sedadv_ax':
d['hatc1']['ax']['sedadv'] = c[:, :, :, i]
elif submod == 'erosion':
d['hatc1']['a']['erosion'] = {}
for j, subsubmod in enumerate(keysu1):
d['hatc1']['a']['erosion'][subsubmod] = c[:, :, :, len(self.submodulesToRun)-1+j]
else:
d['hatc1']['a'][submod] = c[:, :, :, i]
if 'sedadv' not in self.submodulesToRun:
d['hatc1']['ax'] = 0
return d
def svarFunction(Kv, F, Fsurf, Fbed, data, hasMatrix=False):
"""Solve a function
Ds - (Kv s_z)_z = -u0_z + F
subject to
Kv s_z (-H) = Fbed
Kv s_z (0) = Fsurf
The returned solution has a part 'sCoef' and 'szCoef' for the forcing by u_z and a part 'sForced' and 'szForced' for the
forcing by F, Fbed and Fsurf.
Args:
Kv: (ndarray(jmax+1, kmax+1, fmax+1)) - data on eddy diffusivity
or a salinityMatrix as calculated before by this function.
F: (ndarray(jmax+1, kmax+1, fmax+1, nRHS)) - interior forcing. nRHS is the number of individual forcing components
Fsurf: (ndarray(jmax+1, 1, fmax+1, nRHS)) - surface forcing. nRHS is the number of individual forcing components
Fbed: (ndarray(jmax+1, 1, fmax+1, nRHS)) - surface forcing. nRHS is the number of individual forcing components
data: (DataContainer) - should at least contain 'grid', 'OMEGA' and 'u0'
hasMatrix: (bool) - if True then it is assumed that Kv contains a salinityMatrix as calculated by this function before.
The matrix is not computed again, which saves time.
Returns:
sCoef and szCoef: (ndarray(jmax+1, kmax+1, fmax+1, 1)) the solution s and its vertical derivative for the forcing by u0_z.
the final dimension '1' denotes a single forcing component.
sForced and szForced: (ndarray(jmax+1, kmax+1, fmax+1, nRHS)) the solution s and its vertical derivative for the other forcings.
the solution is separated for each RHS term in F, Fsurf and Fbed.
salinityMatrix: matrix used in computation. Can be reused for subsequent calls of this function.
"""
# Init
jmax = data.v('grid', 'maxIndex', 'x') # maximum index of x grid (jmax+1 grid points incl. 0)
kmax = data.v('grid', 'maxIndex', 'z') # maximum index of z grid (kmax+1 grid points incl. 0)
fmax = data.v('grid', 'maxIndex', 'f') # maximum index of f grid (fmax+1 grid points incl. 0)
OMEGA = data.v('OMEGA')
uz = data.d('u0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1), dim='z')
u0bed = data.v('u0', range(0, jmax+1), [-1], range(0, fmax+1))
ftot = 2*fmax+1
# Determine bandwidth of eddy viscosity matrix
bandwidth = 0
for n in np.arange(fmax, -1, -1):
if np.any(abs(Kv[:, :, n]) > 0):
bandwidth = max(bandwidth, n)
# Init Ctd
nRHS = F.shape[-1]
salinityMatrix = np.empty([jmax+1, 2*ftot+2*bandwidth+1, ftot*(kmax+1)], dtype=complex)
szCoef = np.zeros([jmax+1, kmax+1, fmax+1, 1], dtype=complex)
szForced = np.zeros([jmax+1, kmax+1, fmax+1, nRHS], dtype=complex)
# build, save and solve the velocity matrices in every water column
for j in range(0, jmax+1):
# dz vectors
dz = (data.v('grid', 'axis', 'z')[0, 1:]-data.v('grid', 'axis', 'z')[0, 0:-1])*data.n('H', j)
dz = dz.reshape(dz.shape[0])
dz_down = dz[:kmax-1].reshape(kmax-1, 1, 1)
dz_up = dz[1:kmax].reshape(kmax-1, 1, 1)
dz_av = 0.5*(dz_down+dz_up)
##### LEFT HAND SIDE #####
if not hasMatrix:
# Init
A = np.zeros([2*ftot+2*bandwidth+1, ftot*(kmax+1)], dtype=complex)
N = np.zeros([kmax+1, 2*bandwidth+1, ftot], dtype=complex)
# Build eddy viscosity matrix blocks
N[:, bandwidth, :] = Kv[j, :, 0].reshape(kmax+1, 1)*np.ones([1, ftot])
for n in range(1, bandwidth+1):
N[:, bandwidth+n, :-n] = 0.5*Kv[j, :, n].reshape(kmax+1, 1)*np.ones([1, ftot])
N[:, bandwidth-n, n:] = 0.5*np.conj(Kv[j, :, n]).reshape(kmax+1, 1)*np.ones([1, ftot])
# Build matrix. Discretisation: central for second derivative, central for first derivative
# NB. can use general numerical schemes as dz < 0
a = -N[:-2, :, :]/dz_down
b = N[1:kmax, :, :]/dz_up+N[1:kmax, :, :]/dz_down
c = -N[2:, :, :]/dz_up
b[:, bandwidth, :] += (np.arange(-fmax, ftot-fmax)*1j*OMEGA).reshape((1, ftot))*dz_av.reshape((kmax-1, 1))
a = np.swapaxes(a, 0, 1)
b = np.swapaxes(b, 0, 1)
c = np.swapaxes(c, 0, 1)
# Build matrix
A[2*ftot:2*ftot+2*bandwidth+1, :-2*ftot] = a.reshape(a.shape[0], a.shape[1]*a.shape[2])
A[2*fmax+1:2*fmax+2*bandwidth+2, ftot:-ftot] = b.reshape(a.shape[0], a.shape[1]*a.shape[2])
A[0:2*bandwidth+1, 2*ftot:] = c.reshape(a.shape[0], a.shape[1]*a.shape[2])
# Boundary conditions
# Surface (k=0)
A[2*fmax+1:2*fmax+2*bandwidth+2, :ftot] = N[0, :, :]
# Bed (k=kmax)
A[2*fmax+1:2*fmax+2*bandwidth+2, -ftot:] = N[-1, :, :]
# save matrix
salinityMatrix[j, Ellipsis] = A[Ellipsis]
bandwidthA = bandwidth+ftot
else:
A = Kv[j, Ellipsis] # if hasMatrix Av replaces the role of the matrix in this equation
bandwidthA = (A.shape[0]-1)/2
##### RIGHT HAND SIDE #####
# Implicit part of the forcing uz0*dz*Sx
umat = np.zeros((kmax+1, ftot), dtype=complex)
umat[1:-1, fmax:ftot] = -uz[j, 1:-1, :]*dz_av.reshape((kmax-1, 1))
uRHS_implicit = umat.reshape((ftot*(kmax+1), 1))
# Forcing from other factors
uRHS = np.zeros([ftot*(kmax+1), nRHS], dtype=complex)
uRHS[fmax:ftot, :] = Fsurf[j, 0, :, :]
uRHS[-ftot+fmax:, :] = Fbed[j, 0, :, :]
F_full = np.concatenate((np.zeros((jmax+1, kmax+1, fmax, nRHS)), F), 2)
uRHS[ftot:-ftot, :] = 0.5*(F_full[j, 2:, :, :]-F_full[j, :-2, :, :]).reshape((kmax-1)*ftot, nRHS)
##### SOLVE #####
sz = solve_banded((bandwidthA, bandwidthA), A, np.concatenate((uRHS_implicit, uRHS), 1), overwrite_ab=True, overwrite_b=True)
sz = sz.reshape(kmax+1, ftot, 1+nRHS)
szCoef[j, :, :, :] = ny.eliminateNegativeFourier(sz[:, :, :1], 1)
szForced[j, :, :, :] = ny.eliminateNegativeFourier(sz[:, :, 1:], 1)
del sz, uRHS_implicit, uRHS
##### INTEGRATION CONSTANT #####
SIGMASAL = data.v('SIGMASAL')
if hasMatrix:
KvBed = data.v('Av', range(0, jmax+1), [kmax-1, kmax], range(0, fmax+1))/SIGMASAL # TODO: shift to turbulence model
else:
KvBed = Kv[:, [kmax-1, kmax], :]
Dinv = np.diag((np.arange(0, fmax+1)*1j*OMEGA))
Dinv[0, 0] = np.inf
Dinv[range(0, fmax+1), range(0, fmax+1)] = Dinv[range(0, fmax+1), range(0, fmax+1)]**(-1)
z = ny.dimensionalAxis(data.slice('grid'), 'z')
dzbed = z[:, [-1], 0] - z[:, [-2], 0]
sbedCoef = -ny.complexAmplitudeProduct(KvBed[:, [-2], :], szCoef[:, [-2], :, :], 2)/dzbed.reshape(dzbed.shape+(1, 1)) - u0bed.reshape(u0bed.shape+(1,))
sbedCoef = np.dot(Dinv, sbedCoef)
sbedCoef = np.rollaxis(sbedCoef, 0, 3)
sbedForced = (Fbed-ny.complexAmplitudeProduct(KvBed[:, [-2], :], szForced[:, [-2], :, :], 2))/dzbed.reshape(dzbed.shape+(1, 1)) + F[:, [-1], :, :]
sbedForced = np.dot(Dinv, sbedForced)
sbedForced = np.rollaxis(sbedForced, 0, 3)
##### INTEGRATION #####
sCoef = ny.integrate(szCoef, 'z', kmax, np.arange(0, kmax+1), data.slice('grid')) + sbedCoef*np.ones((1, kmax+1, 1, 1))
sForced = ny.integrate(szForced, 'z', kmax, np.arange(0, kmax+1), data.slice('grid')) + sbedForced*np.ones((1, kmax+1, 1, 1))
##### CLOSURE FOR THE DEPTH-AVERAGED TIME-AVERAGED SALINITY VARIATION #####
H = data.v('H', range(0, jmax+1)).reshape((jmax+1, 1, 1, 1))
sCoef[:, :, [0], :] -= (ny.integrate(sCoef[:, :, [0], :], 'z', kmax, 0, data.slice('grid'))/H)*np.ones((1, kmax+1, 1, 1))
sForced[:, :, [0], :] -= (ny.integrate(sForced[:, :, [0], :], 'z', kmax, 0, data.slice('grid'))/H)*np.ones((1, kmax+1, 1, 1))
if hasMatrix:
return sCoef, sForced, szCoef, szForced, Kv
else:
return sCoef, sForced, szCoef, szForced, salinityMatrix
def run(self):
"""
Returns:
Dictionary with results. At least contains the variables listed as output in the registry
"""
self.logger.info('Running module HydroFirst')
# Init
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
G = self.input.v('G')
BETA = self.input.v('BETA')
OMEGA = self.input.v('OMEGA')
ftot = 2 * fmax + 1
submodulesToRun = self.input.v('submodules')
# check if the river term should be compensated for by reference level. Only if river is on, non-zero and there is no leading-order contribution
if 'river' in submodulesToRun and self.input.v(
'Q1') != 0 and not np.any(
self.input.v('zeta0', 'river', range(0, jmax + 1), 0, 0)):
RiverReferenceCompensation = 1
else:
RiverReferenceCompensation = 0
################################################################################################################
# velocity as function of water level
################################################################################################################
## LHS terms
# try to get velocityMatrix from the input. If it is not found, proceed to calculate the matrix again
A = self.input.v(
'velocityMatrix'
) # known that this is numeric data, ask without arguments to retrieve full ndarray
velocityMatrix = True
if A is None:
A = self.input.v('Av', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
velocityMatrix = False
else:
A = A
## RHS terms
# Determine number/names of right hand side
submodulesVelocityForcing = [
i for i in ['adv', 'nostress', 'baroc', 'mixing', 'river']
if i in submodulesToRun
]
#submodulesVelocityConversion = [i for i, mod in enumerate(submodulesToRun) if mod in ['adv', 'nostress', 'baroc', 'mixing']]
submodulesVelocityConversion = [
submodulesToRun.index(i) for i in submodulesVelocityForcing
]
nRHS = len(submodulesVelocityForcing)
F = np.zeros([jmax + 1, kmax + 1, ftot, nRHS], dtype=complex)
Fsurf = np.zeros([jmax + 1, 1, ftot, nRHS], dtype=complex)
Fbed = np.zeros([jmax + 1, 1, ftot, nRHS], dtype=complex)
uFirst = np.zeros([jmax + 1, kmax + 1, ftot,
len(submodulesToRun)],
dtype=complex)
uzFirst = np.zeros([jmax + 1, kmax + 1, ftot,
len(submodulesToRun)],
dtype=complex)
u0 = self.input.v('u0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
zeta0 = self.input.v('zeta0', range(0, jmax + 1), [0],
range(0, fmax + 1))
if RiverReferenceCompensation: # for reference level variation
F[:, :, fmax,
submodulesVelocityForcing.index('river')] = -G * self.input.d(
'R', range(0, jmax + 1), dim='x').reshape(
(jmax + 1, 1)) * np.ones((1, kmax + 1))
if 'adv' in submodulesVelocityForcing:
u0x = self.input.d('u0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
u0z = self.input.d('u0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
w0 = self.input.v('w0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
eta = ny.complexAmplitudeProduct(
u0, u0x, 2) + ny.complexAmplitudeProduct(w0, u0z, 2)
eta = np.concatenate((np.zeros([jmax + 1, kmax + 1, fmax]), eta),
2)
F[:, :, :, submodulesVelocityForcing.index('adv')] = -eta
if 'nostress' in submodulesVelocityForcing:
D = (np.arange(0, fmax + 1) * 1j * OMEGA).reshape(
(1, 1, fmax + 1)) * np.ones((jmax + 1, 1, 1))
zeta0x = self.input.d('zeta0',
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='x')
chi = D * u0[:, [0], :] + G * zeta0x
chi = ny.complexAmplitudeProduct(chi, zeta0, 2)
chi = np.concatenate((np.zeros([jmax + 1, 1, fmax]), chi), 2)
Fsurf[:, :, :, submodulesVelocityForcing.index('nostress')] = -chi
if 'baroc' in submodulesVelocityForcing:
sx = self.input.d('s0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
Jsx = -ny.integrate(
sx, 'z', 0, range(0, kmax + 1), self.input.slice('grid')
) # integral from z to 0 has its boundaries inverted and has a minus sign to compensate
Jsx = np.concatenate((np.zeros([jmax + 1, kmax + 1, fmax]), Jsx),
2)
F[:, :, :,
submodulesVelocityForcing.index('baroc')] = -G * BETA * Jsx
if 'mixing' in submodulesVelocityForcing:
u0z = self.input.d('u0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
Av1 = self.input.v('Av1', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
ksi = ny.complexAmplitudeProduct(u0z, Av1, 2)
ksiz = ny.derivative(ksi, 1, self.input.slice('grid'))
ksi = np.concatenate((np.zeros([jmax + 1, kmax + 1, fmax]), ksi),
2)
ksiz = np.concatenate((np.zeros([jmax + 1, kmax + 1, fmax]), ksiz),
2)
F[:, :, :, submodulesVelocityForcing.index('mixing')] = ksiz
Fsurf[:, :, :,
submodulesVelocityForcing.index('mixing')] = -ksi[:, [0],
Ellipsis]
## Removed 14-7-2017 YMD: Roughness1*u0 and Av1*u0z should be equal, so this term cancels
# if self.input.v('BottomBC') in ['PartialSlip']:
# Fbed[:, :, :, submodulesVelocityForcing.index('mixing')] = -ksi[:,[kmax],Ellipsis]
# roughness1 = self.input.v('Roughness1', range(0, jmax+1), [0], range(0, fmax+1))
# if roughness1 is not None:
# ksi = ny.complexAmplitudeProduct(u0[:, [-1], :], roughness1, 2)
# ksi = np.concatenate((np.zeros([jmax+1, 1, fmax]), ksi), 2)
# Fbed[:, :, :, submodulesVelocityForcing.index('mixing')] = ksi
## Solve equation
uCoef, uFirst[:, :, :,
submodulesVelocityConversion], uzCoef, uzFirst[:, :, :,
submodulesVelocityConversion], _ = uFunctionMomentumConservative(
A,
F,
Fsurf,
Fbed,
self.
input,
hasMatrix
=velocityMatrix
)
################################################################################################################
# water level
################################################################################################################
## LHS terms
# try to get zetaMatrix from the input. If it is not found, proceed to calculate the matrix again
B = self.input.v(
'zetaMatrix'
) # known that this is numeric data, ask without arguments to retrieve full ndarray
zetaMatrix = True
if B is None:
zetaMatrix = False
utemp = uCoef.reshape(
uCoef.shape[:2] + (1, ) + uCoef.shape[2:]
) # reshape as the 'f' dimension is not grid conform; move it to a higher dimension
JuCoef = ny.integrate(utemp, 'z', kmax, 0,
self.input.slice('grid'))
JuCoef = JuCoef.reshape(jmax + 1, 1, ftot,
ftot) # reshape back to original grid
B = -G * JuCoef * self.input.v('B', np.arange(
0, jmax + 1)).reshape(jmax + 1, 1, 1, 1)
## RHS terms
# advection, no-stress, baroclinic
utemp = uFirst.reshape(
uFirst.shape[:2] + (1, ) + uFirst.shape[2:]
) # reshape as the 'f' dimension is not grid conform; move it to a higher dimension
JuFirst = ny.integrate(utemp, 'z', kmax, 0, self.input.slice('grid'))
JuFirst = JuFirst.reshape(
jmax + 1, 1, ftot,
uFirst.shape[-1]) # reshape back to original grid
# stokes
if 'stokes' in submodulesToRun:
gamma = ny.complexAmplitudeProduct(u0[:, 0, None, Ellipsis], zeta0,
2)
gamma = np.concatenate((np.zeros([jmax + 1, 1, fmax]), gamma), 2)
JuFirst[:, :, :, submodulesToRun.index('stokes')] = gamma
BJuFirst = JuFirst * self.input.v('B', np.arange(0, jmax + 1)).reshape(
jmax + 1, 1, 1, 1)
# open BC: tide
Fopen = np.zeros([1, 1, ftot, len(submodulesToRun)], dtype=complex)
if 'tide' in submodulesToRun:
Fopen[0, 0, fmax:,
submodulesToRun.index('tide')] = ny.amp_phase_input(
self.input.v('A1'), self.input.v('phase1'), (fmax + 1, ))
# closed BC: river
Fclosed = np.zeros([1, 1, ftot, len(submodulesToRun)], dtype=complex)
if 'river' in submodulesToRun:
Fclosed[0, 0, fmax,
submodulesToRun.index('river')] = -self.input.v('Q1')
# closed BC: other terms
Fclosed += -JuFirst[jmax, 0, :, :] * self.input.v('B', jmax)
## Solve equation
zetaCoef, zetaxCoef, _ = zetaFunctionMassConservative(
B, BJuFirst, Fopen, Fclosed, self.input, hasMatrix=zetaMatrix)
zetax = ny.eliminateNegativeFourier(zetaxCoef, 2)
zeta = ny.eliminateNegativeFourier(zetaCoef, 2)
################################################################################################################
# velocity
################################################################################################################
u = np.empty((jmax + 1, kmax + 1, ftot, len(submodulesToRun)),
dtype=uCoef.dtype)
uz = np.empty((jmax + 1, kmax + 1, ftot, len(submodulesToRun)),
dtype=uCoef.dtype)
for j in range(0, jmax + 1):
u[j, :, :, :] = np.dot(uCoef[j, :, :, :],
-G * zetaxCoef[j, 0, :, :])
uz[j, :, :, :] = np.dot(uzCoef[j, :, :, :],
-G * zetaxCoef[j, 0, :, :])
u += uFirst
uz += uzFirst
u = ny.eliminateNegativeFourier(u, 2)
uz = ny.eliminateNegativeFourier(uz, 2)
################################################################################################################
# vertical velocity
################################################################################################################
w = self.verticalVelocity(u)
################################################################################################################
# Make final dictionary to return
################################################################################################################
d = {}
d['zeta1'] = {}
d['u1'] = {}
d['w1'] = {}
for i, submod in enumerate(submodulesToRun):
nf = ny.functionTemplates.NumericalFunctionWrapper(
zeta[:, :, :, i], self.input.slice('grid'))
nf.addDerivative(zetax[:, :, :, i], 'x')
d['zeta1'][submod] = nf.function
nfu = ny.functionTemplates.NumericalFunctionWrapper(
u[:, :, :, i], self.input.slice('grid'))
nfu.addDerivative(uz[:, :, :, i], 'z')
d['u1'][submod] = nfu.function
d['w1'][submod] = w[:, :, :, i]
return d
def run(self):
self.logger.info('Running module SalinityLead')
# Init
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
SIGMASAL = self.input.v('SIGMASAL')
################################################################################################################
# First-order salinity variation as function of leading-order salinity closure
################################################################################################################
# build, save and solve the velocity matrices in every water column
Kv = self.input.v('Av', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))/SIGMASAL
F = np.zeros([jmax+1, kmax+1, fmax+1, 0])
Fsurf = np.zeros([jmax+1, 1, fmax+1, 0])
Fbed = np.zeros([jmax+1, 1, fmax+1, 0])
sCoef, _, szCoef, _, salinityMatrix = svarFunction(Kv, F, Fsurf, Fbed, self.input)
################################################################################################################
# Leading-order salinity closure
################################################################################################################
## LHS terms
# First-order river discharge
Q = -self.input.v('Q1', range(0, jmax+1))
# First-order river discharge or zero if not available
# u1riv = self.input.v('u1', 'river', range(0, jmax+1), range(0, kmax+1))
# B = self.input.v('B', range(0, jmax+1))
# if u1riv is None:
# Q = np.zeros((jmax+1))
# self.logger.warning('No first-order river discharge found in module SalinityLead')
# else:
# Q = ny.integrate(u1riv, 'z', kmax, 0, self.input.slice('grid'))
# Q = Q*B
# del u1riv
# Diffusion coefficient
# part 1) diffusive transport
H = self.input.v('H', range(0, jmax+1))
B = self.input.v('B', range(0, jmax+1))
Kh = self.input.v('Kh', range(0, jmax+1))
AK = np.real(H*B*Kh)
# part 2) advective transport
u0 = self.input.v('u0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
us = ny.complexAmplitudeProduct(u0, sCoef, 2)[:, :, 0, 0] # subtidal part of u*s
us = ny.integrate(us, 'z', kmax, 0, self.input.slice('grid')).reshape(jmax+1)
AK += -np.real(B*us)
del us, u0, H, B
## RHS terms
F = np.zeros([jmax+1, 1]) # dimensions (jmax+1, NumberOfForcings)
# open BC: effect of the sea
Fopen = np.zeros([1, 1])
Fopen[0, 0] = self.input.v('ssea')
# closed BC: always zero. This is by assumption that the salinity vanishes for x=L
Fclosed = np.zeros([1, 1])
## Solve equation
S0 = np.zeros((jmax+1, 1, fmax+1, 1))
Sx0 = np.zeros((jmax+1, 1, fmax+1, 1))
S0[:, 0, 0, :], Sx0[:, 0, 0, :] = sclosureFunction((Q, AK), F, Fopen, Fclosed, self.input)
################################################################################################################
# First-order salinity variation
################################################################################################################
s1 = ny.complexAmplitudeProduct(sCoef, Sx0, 2)
sz1 = ny.complexAmplitudeProduct(szCoef, Sx0, 2)
################################################################################################################
# Make final dictionary to return
################################################################################################################
d = {}
d['salinityMatrix'] = salinityMatrix
nf = ny.functionTemplates.NumericalFunctionWrapper(S0[:, :, :, 0], self.input.slice('grid'))
nf.addDerivative(Sx0[:, :, :, 0], 'x')
d['s0'] = nf.function
nf2 = ny.functionTemplates.NumericalFunctionWrapper(s1[:, :, :, 0], self.input.slice('grid'))
nf2.addDerivative(sz1[:, :, :, 0], 'z')
d['s1var'] = nf2.function
return d
def run(self):
self.logger.info('Running module StaticAvailability')
################################################################################################################
## Init
################################################################################################################
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
c0 = self.input.v('hatc0', 'a', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
c1_a0 = self.input.v('hatc1', 'a', range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
c1_a0x = self.input.v('hatc1', 'ax', range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
if isinstance(c1_a0x, bool):
c1_a0x = np.zeros((jmax + 1, kmax + 1, fmax + 1))
d = {}
c0_int = ny.integrate(c0, 'z', kmax, 0, self.input.slice('grid'))
B = self.input.v('B', range(0, jmax + 1), [0], [0])
u0 = self.input.v('u0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
zeta0 = self.input.v('zeta0', range(0, jmax + 1), [0],
range(0, fmax + 1))
Kh = self.input.v('Kh', range(0, jmax + 1), [0], [0])
################################################################################################################
## Second order closure
################################################################################################################
u1 = self.input.v('u1', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
d['T'] = {}
d['F'] = {}
T0 = 0
F0 = 0
## Transport T ############################################################################################
## T.1. - u0*c1_a0
# Total
c1a_f0 = c1_a0
T0 += ny.integrate(ny.complexAmplitudeProduct(u0, c1a_f0, 2), 'z',
kmax, 0, self.input.slice('grid'))
# Decomposition
for submod in self.input.getKeysOf('hatc1', 'a'):
if submod == 'erosion':
for subsubmod in self.input.getKeysOf('hatc1', 'a', 'erosion'):
c1_a0_comp = self.input.v('hatc1', 'a', submod, subsubmod,
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1))
c1a_f0_comp_res = c1_a0_comp
d['T'] = self.dictExpand(
d['T'], subsubmod,
['TM' + str(2 * n) for n in range(0, fmax + 1)
]) # add submod index to dict if not already
# transport with residual availability
for n in range(0, fmax + 1):
tmp = np.zeros(c1a_f0_comp_res.shape, dtype=complex)
tmp[:, :, n] = c1a_f0_comp_res[:, :, n]
tmp = ny.integrate(
ny.complexAmplitudeProduct(u0, tmp, 2), 'z', kmax,
0, self.input.slice('grid'))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T'][subsubmod]['TM' + str(2 * n)] += tmp
else:
c1_a0_comp = self.input.v('hatc1', 'a', submod,
range(0,
jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
c1a_f0_comp_res = c1_a0_comp
d['T'] = self.dictExpand(
d['T'], submod,
['TM' + str(2 * n) for n in range(0, fmax + 1)
]) # add submod index to dict if not already
# transport with residual availability
for n in range(0, fmax + 1):
tmp = np.zeros(c1a_f0_comp_res.shape, dtype=complex)
tmp[:, :, n] = c1a_f0_comp_res[:, :, n]
tmp = ny.integrate(ny.complexAmplitudeProduct(u0, tmp, 2),
'z', kmax, 0,
self.input.slice('grid'))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T'][submod]['TM' + str(2 * n)] += tmp
## T.2. - u1*c0
# Total
T0 += ny.integrate(ny.complexAmplitudeProduct(u1, c0, 2), 'z', kmax, 0,
self.input.slice('grid'))
# Decomposition
for submod in self.input.getKeysOf('u1'):
u1_comp = self.input.v('u1', submod, range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
d['T'] = self.dictExpand(
d['T'], submod,
['TM' + str(2 * n) for n in range(0, fmax + 1)
]) # add submod index to dict if not already
# transport with residual availability
for n in range(0, fmax + 1):
tmp = np.zeros(u1_comp.shape, dtype=complex)
tmp[:, :, n] = u1_comp[:, :, n]
if submod == 'stokes':
tmp = ny.integrate(ny.complexAmplitudeProduct(tmp, c0, 2),
'z', kmax, 0,
self.input.slice('grid'))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T'][submod] = self.dictExpand(
d['T'][submod], 'TM' + str(2 * n),
['return', 'drift'])
d['T'][submod]['TM0']['return'] += tmp
else:
tmp = ny.integrate(ny.complexAmplitudeProduct(tmp, c0, 2),
'z', kmax, 0,
self.input.slice('grid'))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T'][submod]['TM' + str(2 * n)] += tmp
## T.5. - u0*c0*zeta0
# Total
T0 += ny.complexAmplitudeProduct(
ny.complexAmplitudeProduct(u0[:, [0], :], c0[:, [0], :], 2), zeta0,
2)
# Decomposition
uzeta = ny.complexAmplitudeProduct(u0[:, [0], :], zeta0, 2)
d['T'] = self.dictExpand(
d['T'], 'stokes', ['TM' + str(2 * n) for n in range(0, fmax + 1)])
# transport with residual availability
for n in range(0, fmax + 1):
tmp = np.zeros(c0[:, [0], :].shape, dtype=complex)
tmp[:, :, n] = c0[:, [0], n]
tmp = ny.complexAmplitudeProduct(uzeta, tmp, 2)[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T']['stokes']['TM' + str(2 * n)]['drift'] += tmp
## T.6. - u1riv*c2rivriv
c2 = self.input.v('hatc2', 'a', 'erosion', 'river_river',
range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
u1riv = self.input.v('u1', 'river', range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
if u1riv is not None:
d['T'] = self.dictExpand(
d['T'], 'river_river',
'TM0') # add submod index to dict if not already
tmp = ny.integrate(ny.complexAmplitudeProduct(u1riv, c2, 2), 'z',
kmax, 0, self.input.slice('grid'))
if any(abs(tmp[:, 0, 0])) > 10**-14:
d['T']['river_river']['TM0'] = tmp[:, 0, 0]
T0 += tmp
## T.7. - diffusive part
# Total
c0x = self.input.d('hatc0',
'a',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
T0 += -Kh * ny.integrate(c0x, 'z', kmax, 0, self.input.slice('grid'))
c2x = self.input.d('hatc2',
'a',
'erosion',
'river_river',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
T0 += -Kh * ny.integrate(c2x, 'z', kmax, 0, self.input.slice('grid'))
# Decomposition
d['T'] = self.dictExpand(d['T'], 'diffusion_tide', ['TM0'])
d['T'] = self.dictExpand(d['T'], 'diffusion_river', ['TM0'])
# transport with residual availability
tmp = -(Kh * ny.integrate(c0x, 'z', kmax, 0,
self.input.slice('grid')))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T']['diffusion_tide']['TM0'] = tmp
tmp = -(Kh * ny.integrate(c2x, 'z', kmax, 0,
self.input.slice('grid')))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['T']['diffusion_river']['TM0'] = tmp
## Diffusion F ############################################################################################
## F.1. - u0*C1ax*f0
# Total
F0 += ny.integrate(ny.complexAmplitudeProduct(u0, c1_a0x, 2), 'z',
kmax, 0, self.input.slice('grid'))
# Decomposition
for submod in self.input.getKeysOf('hatc1', 'ax'):
c1_ax0_comp = self.input.v('hatc1', 'ax', submod,
range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
d['F'] = self.dictExpand(
d['F'], submod,
['FM' + str(2 * n) for n in range(0, fmax + 1)
]) # add submod index to dict if not already
# transport with residual availability
for n in range(0, fmax + 1):
tmp = np.zeros(u0.shape, dtype=complex)
tmp[:, :, n] = u0[:, :, n]
tmp = ny.integrate(
ny.complexAmplitudeProduct(tmp, c1_ax0_comp, 2), 'z', kmax,
0, self.input.slice('grid'))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['F'][submod]['FM' + str(2 * n)] += tmp
## F.3. - diffusive part
# Total
F0 += -Kh * ny.integrate(c0, 'z', kmax, 0, self.input.slice('grid'))
F0 += -Kh * ny.integrate(c2, 'z', kmax, 0, self.input.slice('grid'))
# Decomposition
d['F'] = self.dictExpand(d['F'], 'diffusion_tide', ['FM0'])
d['F'] = self.dictExpand(d['F'], 'diffusion_river', ['FM0'])
# transport with residual availability
tmp = -(Kh * ny.integrate(c0, 'z', kmax, 0,
self.input.slice('grid')))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['F']['diffusion_tide']['FM0'] = tmp
tmp = -(Kh * ny.integrate(c2, 'z', kmax, 0,
self.input.slice('grid')))[:, 0, 0]
if any(abs(tmp)) > 10**-14:
d['F']['diffusion_river']['FM0'] = tmp
## Solve ################################################################################################
## Add all mechanisms & compute a0c
from src.DataContainer import DataContainer
dc = DataContainer(d)
dc.merge(self.input.slice('grid'))
T_til = np.real(dc.v('T', range(0, jmax + 1)))
F_til = np.real(dc.v('F', range(0, jmax + 1)))
# DEBUG: CHECKS IF COMPOSITE T, F == total T, F
# print np.max(abs((dc.v('T', range(0, jmax+1))-T0[:, 0, 0])/(T0[:, 0, 0]+10**-10)))
# print np.max(abs((dc.v('F', range(0, jmax+1))-F0[:, 0, 0])/(F0[:, 0, 0]+10**-10)))
integral = -ny.integrate(T_til / (F_til - 10**-6), 'x', 0,
range(0, jmax + 1), self.input.slice('grid'))
if self.input.v('Qsed') is None:
G = 0
else:
G = self.input.v('Qsed') / B[-1, 0, 0]
P = ny.integrate(G / (F_til - 10**-6) * np.exp(-integral), 'x', 0,
range(0, jmax + 1), self.input.slice('grid'))
################################################################################################################
# Boundary condition 1
################################################################################################################
if self.input.v('sedbc') == 'astar':
astar = self.input.v('astar')
k = astar * ny.integrate(B[:, 0, 0], 'x', 0, jmax,
self.input.slice('grid')) / ny.integrate(
B[:, 0, 0] * np.exp(integral), 'x', 0,
jmax, self.input.slice('grid'))
f0 = (k - P) * np.exp(integral)
f0x = (-T_til * f0 - G) / (F_til - 10**-6)
################################################################################################################
# Boundary condition 2
################################################################################################################
elif self.input.v('sedbc') == 'csea':
csea = self.input.v('csea')
c000 = np.real(c0_int[0, 0, 0])
k = csea / c000 * (self.input.v('grid', 'low', 'z', 0) -
self.input.v('grid', 'high', 'z', 0))
f0 = (k - P) * np.exp(integral)
f0x = (-T_til * f0 - G) / (F_til - 10**-6)
else:
from src.util.diagnostics.KnownError import KnownError
raise KnownError(
'sediment boundary sedbc not known: use astar or csea')
################################################################################################################
# Store in dict
################################################################################################################
d['a'] = f0
d['c0'] = c0 * f0.reshape((jmax + 1, 1, 1))
d['c1'] = c1_a0 * f0.reshape((jmax + 1, 1, 1)) + c1_a0x * f0x.reshape(
(jmax + 1, 1, 1))
d['c2'] = c2 * f0.reshape((jmax + 1, 1, 1))
return d
def uRelax(self, order, init):
"""Compute the absolute velocity and absolute velocity times the depth at 'order',
i.e. |u|^<order>, (|u|(H+R+zeta))^<order>.
Then make a relaxation of these signals using the the previous iteration and relaxtion factor set as class var.
Implements two methods:
order == None: truncation method. Else: scaling method
init (bool): initial iteration?
"""
# Init
jmax = self.input.v('grid', 'maxIndex', 'x') # maximum index of x grid (jmax+1 grid points incl. 0)
kmax = self.input.v('grid', 'maxIndex', 'z') # maximum index of z grid (kmax+1 grid points incl. 0)
fmax = self.input.v('grid', 'maxIndex', 'f') # maximum index of f grid (fmax+1 grid points incl. 0)
depth = self.input.v('grid', 'low', 'z', range(0, jmax+1), [0], [0]) - self.input.v('grid', 'high', 'z', range(0, jmax+1), [0], [0])
# test if present u is grid-conform (may not be if the previous runs were on a different resolution)
utest = self.input.v('u0', range(0, jmax+1), range(0, kmax+1), range(0, fmax+1))
no_u = False
if isinstance(utest, bool):
no_u = True
################################################################################################################
# 1. make the absolute velocity
################################################################################################################
c = ny.polyApproximation(np.abs, 8) # chebyshev coefficients for abs
## Truncation method
if order == None:
## 1a. Gather velocity and zeta components
zeta = 0
u = 0
comp = 0
while self.input.v('zeta'+str(comp)) and self.input.v('u'+str(comp)) and comp <= self.truncationorder and not no_u:
zeta += self.input.v('zeta'+str(comp), range(0, jmax + 1), [0], range(0, fmax + 1))
u += self.input.v('u'+str(comp), range(0, jmax + 1), range(0, kmax + 1), range(0, fmax + 1))
for submod in self.ignoreSubmodule: # remove submodules to be ignored
try:
zeta -= self.input.v('zeta'+str(comp), submod, range(0, jmax + 1), [0], range(0, fmax + 1))
u -= self.input.v('u'+str(comp), submod, range(0, jmax + 1), range(0, kmax + 1), range(0, fmax + 1))
except:
pass
comp += 1
# if no data for u and zeta is in de DC, then take an initial estimate
if comp == 0:
u = np.zeros((jmax+1, kmax+1, fmax+1))
u[:, :, 0] = 1.
zeta = np.zeros((jmax+1, 1, fmax+1))
usurf = u[:, [0], :]
u = ny.integrate(u, 'z', kmax, 0, self.input.slice('grid'))
## 1b. Divide velocity by a maximum amplitude
uamp = [(np.sum(np.abs(u), axis=-1)+10**-3).reshape((jmax+1, 1, 1)), (np.sum(np.abs(usurf), axis=-1)+10**-3).reshape((jmax+1, 1, 1))]
u = u/uamp[0]
usurf = usurf/uamp[1]
## 1c. Make absolute value
uabs = [np.zeros(u.shape, dtype=complex), np.zeros(u.shape, dtype=complex)] # uabs at depth-av, surface
for n in [0, 1]: # compute for DA (0) and surface (1)
if n == 0:
ut = u
else:
ut = usurf
uabs[n][:, :, 0] = c[0]
u2 = ny.complexAmplitudeProduct(ut, ut, 2)
uabs[n] += c[2]*u2
u4 = ny.complexAmplitudeProduct(u2, u2, 2)
uabs[n] += c[4]*u4
u6 = ny.complexAmplitudeProduct(u2, u4, 2)
uabs[n] += c[6]*u6
del u2, u6
u8 = ny.complexAmplitudeProduct(u4, u4, 2)
uabs[n] += c[8]*u8
del u4, u8
uabs[n] = uabs[n] * uamp[n].reshape((jmax+1, 1, 1))
# Absolute velocity * depth
uabsH = uabs[0] + ny.complexAmplitudeProduct(uabs[1], zeta, 2)
uabs = uabs[0]/depth + ny.complexAmplitudeProduct(uabs[1]-uabs[0]/depth, zeta, 2)
## Scaling method
else:
## 1a. Gather velocity and zeta components
zeta = []
u = []
usurf = []
for comp in range(0, order+1):
if self.input.v('zeta'+str(comp)) and self.input.v('u'+str(comp)) and not no_u:
zetatemp = self.input.v('zeta'+str(comp), range(0, jmax + 1), [0], range(0, fmax + 1))
utemp = self.input.v('u'+str(comp), range(0, jmax + 1), range(0, kmax + 1), range(0, fmax + 1))
for submod in self.ignoreSubmodule: # remove submodules to be ignored
try:
zetatemp -= self.input.v('zeta'+str(comp), submod, range(0, jmax + 1), [0], range(0, fmax + 1))
utemp -= self.input.v('u'+str(comp), submod, range(0, jmax + 1), range(0, kmax + 1), range(0, fmax + 1))
except:
pass
# if no u and zeta in DC, then ..
elif comp == 0: # .. add velocity 1 to subtidal leading order
zetatemp = np.zeros((jmax+1, 1, fmax+1))
utemp = np.zeros((jmax+1, kmax+1, fmax+1))
utemp[:, :, 0] = 1.
else: # .. add nothing at higher orders
zetatemp = np.zeros((jmax+1, 1, fmax+1))
utemp = np.zeros((jmax+1, kmax+1, fmax+1))
zeta.append(zetatemp)
usurf.append(utemp[:, [0], :])
u.append(ny.integrate(utemp, 'z', kmax, 0, self.input.slice('grid')) / depth)
## 1b. Divide velocity by a maximum amplitude
uamp = []
uamp.append((np.sum(np.abs(sum(u)), axis=-1)+10**-3).reshape((jmax+1, 1, 1)))
uamp.append((np.sum(np.abs(sum(usurf)), axis=-1)+10**-3).reshape((jmax+1, 1, 1)))
u = [i/uamp[0] for i in u]
usurf = [i/uamp[1] for i in usurf]
## 1c. Make absolute value
uabs = [np.zeros(u[0].shape+(order+1,), dtype=complex), np.zeros(u[0].shape+(order+1,), dtype=complex)]
for n in [0, 1]: # compute for DA (0) and surface (1)
if n == 0:
ut = u
else:
ut = usurf
uabs[n][:, :, 0, 0] = c[0]
for q in range(0, order+1):
uabs[n][:, :, :, q] += c[2]*self.umultiply(2, q, ut)
uabs[n][:, :, :, q] += c[4]*self.umultiply(4, q, ut)
uabs[n][:, :, :, q] += c[6]*self.umultiply(6, q, ut)
uabs[n][:, :, :, q] += c[8]*self.umultiply(8, q, ut)
uabs[n] = uabs[n] * uamp[n].reshape((jmax+1, 1, 1, 1))
# Absolute velocity * depth
uabsH = uabs[0][:, :, :, order]*depth
for q in range(0, order):
uabsH += ny.complexAmplitudeProduct(uabs[1][:, :, :, q]-uabs[0][:, :, :, q], zeta[order-q-1], 2)
# Only keep uabs at current order
uabs = uabs[0][:, :, :, order] # only keep DA part
################################################################################################################
# 2. Relaxtion
################################################################################################################
## 2a. Relaxation on uabs
if hasattr(self, 'u_prev_iter') and self.u_prev_iter.shape == uabs.shape:
u_prev_iter = self.u_prev_iter
else:
u_prev_iter = uabs # take current uabs if no previous is available
u_prev_iter2 = u_prev_iter - (uabs)
u0 = np.max((uabs, np.min((u_prev_iter2 * (1 - self.RELAX), u_prev_iter2 * (1. + self.RELAX)), axis=0) + (uabs)), axis=0)
u0 = np.min((u0, np.max((u_prev_iter2 * (1 - self.RELAX), u_prev_iter2 * (1. + self.RELAX)), axis=0) + uabs), axis=0)
self.u_prev_iter = u0 # save velocity at bed for next iteration
## 2b. Relaxation on uabs*depth
if hasattr(self, 'uH_prev_iter') and self.uH_prev_iter.shape == uabsH.shape:
u_prev_iter = self.uH_prev_iter
else:
u_prev_iter = uabsH # take current uabsH if no previous is available
u_prev_iter2 = u_prev_iter - (uabsH)
uH0 = np.max((uabsH, np.min((u_prev_iter2 * (1 - self.RELAX), u_prev_iter2 * (1. + self.RELAX)), axis=0) + (uabsH)), axis=0)
uH0 = np.min((uH0, np.max((u_prev_iter2 * (1 - self.RELAX), u_prev_iter2 * (1. + self.RELAX)), axis=0) + uabsH), axis=0)
self.uH_prev_iter = uH0 # save velocity at bed for next iteration
return u0, uH0
def first_order(self, d):
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
ftot = 2 * fmax + 1
OMEGA = self.input.v('OMEGA')
################################################################################################################
# Forcing terms
################################################################################################################
if 'sedadv' in self.submodulesToRun and 'sedadv_ax' not in self.submodulesToRun:
self.submodulesToRun.append('sedadv_ax')
# determine number of submodules
nRHS = len(self.submodulesToRun)
if 'erosion' in self.submodulesToRun:
keysu1 = self.input.getKeysOf('u1')
try:
keysu1.remove(
'mixing'
) # flow due to 'mixing' should not be included in the erosion term
except:
pass
self.submodulesToRun.remove(
'erosion') # move to the end of the list
self.submodulesToRun.append('erosion')
nRHS = nRHS - 1 + len(self.input.getKeysOf('u1'))
F = np.zeros([jmax + 1, kmax + 1, ftot, nRHS], dtype=complex)
Fsurf = np.zeros([jmax + 1, 1, ftot, nRHS], dtype=complex)
Fbed = np.zeros([jmax + 1, 1, ftot, nRHS], dtype=complex)
self.input.merge(d)
c0 = self.input.v('hatc0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
cx0 = self.input.d('hatc0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
cz0 = self.input.d('hatc0',
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
# 1. Erosion
if 'erosion' in self.submodulesToRun:
# erosion due to first-order bed shear stress
# E = erosion(self.ws, Av, 1, self.input, self.erosion_method) # 24-04-2017 Obsolete
# Fbed[:, :, fmax:, self.submodulesToRun.index('erosion')] = -E # 24-04-2017 Obsolete
for submod in keysu1:
E = erosion(self.ws,
1,
self.input,
self.erosion_method,
submodule=(None, submod),
friction=self.frictionpar)
Fbed[:, :, fmax:,
len(self.submodulesToRun) - 1 + keysu1.index(submod)] = -E
# 2. Advection
if 'sedadv' in self.submodulesToRun:
u0 = self.input.v('u0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
w0 = self.input.v('w0', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
eta = ny.complexAmplitudeProduct(
u0, cx0, 2) + ny.complexAmplitudeProduct(w0, cz0, 2)
F[:, :, fmax:, self.submodulesToRun.index('sedadv')] = -eta
F[:, :, fmax:,
self.submodulesToRun.
index('sedadv_ax')] = -ny.complexAmplitudeProduct(u0, c0, 2)
# 3. First-order fall velocity
if 'fallvel' in self.submodulesToRun:
# surface and internal terms
ws1 = self.input.v('ws1', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
ksi = ny.complexAmplitudeProduct(ws1, c0, 2)
ksiz = ny.derivative(ksi, 'z', self.input.slice('grid'))
F[:, :, fmax:, self.submodulesToRun.index('fallvel')] = ksiz
Fsurf[:, 0, fmax:,
self.submodulesToRun.index('fallvel')] = -ksi[:, 0, :]
# adjustment to erosion; only if erosion depends on the settling velocity
if self.erosion_method == 'Chernetsky':
E = erosion(ws1,
0,
self.input,
self.erosion_method,
friction=self.frictionpar)
Fbed[:, :, fmax:, self.submodulesToRun.index('fallvel')] = -E
# 4. First-order eddy diffusivity
if 'mixing' in self.submodulesToRun:
# surface, bed and internal terms
Kv1 = self.input.v('Kv1', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
psi = ny.complexAmplitudeProduct(Kv1, cz0, 2)
psiz = ny.derivative(psi, 'z', self.input.slice('grid'))
F[:, :, fmax:, self.submodulesToRun.index('mixing')] = psiz
Fsurf[:, 0, fmax:,
self.submodulesToRun.index('mixing')] = -psi[:, 0, :]
Fbed[:, 0, fmax:,
self.submodulesToRun.index('mixing')] = -psi[:, -1, :]
# adjustment to erosion
E = erosion(self.ws,
1,
self.input,
self.erosion_method,
submodule=(None, 'mixing'),
friction=self.frictionpar)
Fbed[:, :, fmax:, self.submodulesToRun.index('mixing')] = -E
# 5. No-flux surface correction
if 'noflux' in self.submodulesToRun:
zeta0 = self.input.v('zeta0', range(0, jmax + 1), [0],
range(0, fmax + 1))
D = np.zeros((jmax + 1, 1, fmax + 1, fmax + 1), dtype=complex)
D[:, :, range(0, fmax + 1),
range(0, fmax + 1)] = np.arange(0, fmax + 1) * 1j * OMEGA
Dc0 = ny.arraydot(D, c0[:, [0], :])
chi = ny.complexAmplitudeProduct(Dc0, zeta0, 2)
Fsurf[:, :, fmax:, self.submodulesToRun.index('noflux')] = -chi
################################################################################################################
# Solve equation
################################################################################################################
cmatrix = self.input.v('cMatrix')
if cmatrix is not None:
c, cMatrix = cFunction(None,
cmatrix,
F,
Fsurf,
Fbed,
self.input,
hasMatrix=True)
else:
c, cMatrix = cFunction(self.ws,
self.Kv,
F,
Fsurf,
Fbed,
self.input,
hasMatrix=False)
c = c.reshape((jmax + 1, kmax + 1, ftot, nRHS))
c = ny.eliminateNegativeFourier(c, 2)
################################################################################################################
# Prepare output
################################################################################################################
d['hatc1'] = {}
d['hatc1']['a'] = {}
d['hatc1']['ax'] = {}
for i, submod in enumerate(self.submodulesToRun):
if submod == 'sedadv_ax':
d['hatc1']['ax']['sedadv'] = c[:, :, :, i]
elif submod == 'erosion':
d['hatc1']['a']['erosion'] = {}
for j, subsubmod in enumerate(keysu1):
d['hatc1']['a']['erosion'][
subsubmod] = c[:, :, :,
len(self.submodulesToRun) - 1 + j]
else:
d['hatc1']['a'][submod] = c[:, :, :, i]
if 'sedadv' not in self.submodulesToRun:
d['hatc1']['ax'] = 0
return d
def run(self):
self.currentOrder = self.input.v('order')
if self.currentOrder < 2:
return
# Init
if self.currentOrder == 2:
maxOrder = self.input.v('maxOrder')
jmax = self.input.v('grid', 'maxIndex', 'x')
fmax = self.input.v('grid', 'maxIndex', 'f')
OMEGA = self.input.v('OMEGA')
G = self.input.v('G')
self.submodulesToRun = self.input.v('submodules')
# initialise arrays with surface stress and velocity data
self.surf_stress = np.nan * np.empty(
(jmax + 1, 1, fmax + 1, maxOrder, maxOrder),
dtype=complex) # x, z, f, order, derivative
self.surf_u_der = np.nan * np.empty(
(jmax + 1, 1, fmax + 1, maxOrder, maxOrder + 2),
dtype=complex) # x, z, f, order, derivative
# update those parts of these arrays that are not updated later
# surface stress
D = (np.arange(0, fmax + 1) * 1j * OMEGA).reshape(
(1, 1, fmax + 1)) * np.ones((jmax + 1, 1, 1))
u0 = self.input.v('u0', range(0, jmax + 1), [0],
range(0, fmax + 1))
zeta0x = self.input.d('zeta0',
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='x')
self.surf_stress[:, :, :, 0, 0] = D * u0[:, [0], :] + G * zeta0x
# surface der of u
self.surf_u_der[:, :, :, 0, 0] = u0
self.surf_u_der[:, :, :, 0, 1] = self.input.d('u0',
range(0, jmax + 1),
[0],
range(0, fmax + 1),
dim='z')
Av = self.input.v('Av', range(0, jmax + 1), [0],
range(0, fmax + 1))
Avz = self.input.d('Av',
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='z')
u0z = self.input.d('u0',
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='z')
self.surf_u_der[:, :, :, 0, 2] = -self.Avinv_multiply(
Av, (ny.complexAmplitudeProduct(Avz, u0z, 2) -
self.surf_stress[:, :, :, 0, 0]))
self.logger.info('Running module HydroHigher - order ' +
str(self.currentOrder))
# Init
jmax = self.input.v('grid', 'maxIndex', 'x')
kmax = self.input.v('grid', 'maxIndex', 'z')
fmax = self.input.v('grid', 'maxIndex', 'f')
G = self.input.v('G')
BETA = self.input.v('BETA')
ftot = 2 * fmax + 1
try:
maxContributions = int(self.input.v('maxContributions'))
except:
maxContributions = self.input.v('maxContributions')
d = {}
# update surf_stress and surf_u_der
self.updateSurfaceData()
################################################################################################################
# Velocity in terms of water level gradient
################################################################################################################
## LHS terms
# try to get velocityMatrix from the input. If it is not found, proceed to calculate the matrix again
A = self.input.v(
'velocityMatrix'
) # known that this is numeric data, ask without arguments to retrieve full ndarray
velocityMatrix = True
if A is None:
A = self.input.v('Av', range(0, jmax + 1), range(0, kmax + 1),
range(0, fmax + 1))
velocityMatrix = False
## RHS terms
# Determine number/names of right hand side
nRHS, nRHSVelocity = self.__numberOfForcings()
f_index = -1
f_names = []
# instantiate the forcing components in the equation
F = np.zeros([jmax + 1, kmax + 1, ftot, nRHSVelocity], dtype=complex)
Fsurf = np.zeros([jmax + 1, 1, ftot, nRHSVelocity], dtype=complex)
Fbed = np.zeros([jmax + 1, 1, ftot, nRHSVelocity], dtype=complex)
JuFirst = np.zeros([jmax + 1, 1, ftot, nRHS], dtype=complex)
uFirst = np.zeros([jmax + 1, kmax + 1, ftot, nRHS], dtype=complex)
uzFirst = np.zeros([jmax + 1, kmax + 1, ftot, nRHS], dtype=complex)
# determine RHS terms per submodule - first for velocity
# 1. Advection
if 'adv' in self.submodulesToRun:
for order1 in range(0, self.currentOrder):
order2 = self.currentOrder - order1 - 1
# labels and submodules
u_str1 = 'u' + str(order1)
u_keys1 = self.input.getKeysOf(u_str1)
u_str2 = 'u' + str(order2)
u_keys2 = self.input.getKeysOf(u_str2)
w_str = 'w' + str(order1)
# retrieve data and make advection forcing
for submod1 in u_keys1:
for submod2 in u_keys2:
u = self.input.v(u_str1, submod1, range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1))
ux = self.input.d(u_str2,
submod2,
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
w = self.input.v(w_str, submod1, range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1))
uz = self.input.d(u_str2,
submod2,
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
eta = ny.complexAmplitudeProduct(
u, ux, 2) + ny.complexAmplitudeProduct(w, uz, 2)
eta = np.concatenate(
(np.zeros([jmax + 1, kmax + 1, fmax]), eta), 2)
f_index += 1
f_names.append([
'adv',
submod1 + str(order1) + '-' + submod2 + str(order2)
])
F[:, :, :, f_index] = -eta
del eta
# 2. No-Stress
if 'nostress' in self.submodulesToRun:
chi = 0
for m in range(1, self.currentOrder + 1):
for k in range(0, self.currentOrder - m + 1):
zetapermutations = self.multiindex(
m, self.currentOrder - m - k)
# a. make the zeta^m product
zetasum = 0
for perm in range(0, zetapermutations.shape[0]):
zetaname = 'zeta' + str(zetapermutations[perm, 0])
zeta = self.input.v(zetaname, range(0, jmax + 1), [0],
range(0, fmax + 1))
for comp in range(1, zetapermutations.shape[1]):
zetaname = 'zeta' + str(zetapermutations[perm,
comp])
zeta2 = self.input.v(zetaname, range(0, jmax + 1),
[0], range(0, fmax + 1))
zeta = ny.complexAmplitudeProduct(zeta, zeta2, 2)
zetasum = zetasum + zeta
# b. make the (Av*uz)^(m) term (use m-1 as surf_stress contains (Av*uz)_z^(m))
Avuz = self.surf_stress[:, :, :, k, m - 1]
# c. add all to chi
chi += 1. / np.math.factorial(
m) * ny.complexAmplitudeProduct(Avuz, zetasum, 2)
chi = np.concatenate((np.zeros([jmax + 1, 1, fmax]), chi), 2)
f_index += 1
f_names.append(['nostress', ''])
Fsurf[:, :, :, f_index] = -chi
del chi
# 3. Density-drift # TODO
# if 'densitydrift' in self.submodulesToRun:
# beta_delta = 0
# for m in range(1, self.currentOrder):
# for k in range(0, self.currentOrder-m):
# zetapermutations = self.multiindex(m, self.currentOrder-m-k-1)
# # a. make the zeta^m product
# zetasum = 0
# for perm in range(0, zetapermutations.shape[0]):
# zetaname = 'zeta'+str(zetapermutations[perm, 0])
# zeta = self.input.v(zetaname, range(0, jmax+1), [0], range(0, fmax+1))
# for comp in range(1, zetapermutations.shape[1]):
# zetaname = 'zeta'+str(zetapermutations[perm, comp])
# zeta2 = self.input.v(zetaname, range(0, jmax+1), [0], range(0, fmax+1))
# zeta = ny.complexAmplitudeProduct(zeta, zeta2, 2)
# zetasum = zetasum + zeta
#
# # b. make the (s_x)^(m) term
# sname = 's'+str(k)
# sx = self.input.d(sname, range(0, jmax+1), range(0, kmax+1), range(0, fmax+1), dim='x')
# sxmdz = self.surfaceDerivative(sx, m-1, self.NUMORDER_SURFACE)
#
# # c. add all to beta_delta
# beta_delta += G*BETA*1./np.math.factorial(m)*ny.complexAmplitudeProduct(sxmdz, zetasum, 2)
#
# beta_delta = np.concatenate((np.zeros([jmax+1, 1, fmax]), beta_delta), 2)
#
# f_index += 1
# f_names.append(['densitydrift', ''])
# F[:, :, :, f_index] = -beta_delta*np.ones((jmax+1, kmax+1, fmax+1))
# 4. Baroclinic
# Only use this when salinity on lower order is available
if 'baroc' in self.submodulesToRun:
s_str = 's' + str(self.currentOrder - 1)
sx = self.input.d(s_str,
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='x')
if sx is not None:
Jsx = -G * BETA * ny.integrate(
sx, 'z', 0, range(0, kmax + 1), self.input.slice('grid')
) # integral from z to 0 has its boundaries inverted and has a minus sign to compensate
Jsx = np.concatenate(
(np.zeros([jmax + 1, kmax + 1, fmax]), Jsx), 2)
f_index += 1
f_names.append(['baroc', ''])
F[:, :, :, f_index] = -Jsx
del Jsx, sx
# 5. Mixing
if 'mixing' in self.submodulesToRun:
ksi = np.zeros([jmax + 1, kmax + 1, fmax + 1], dtype=complex)
for m in range(1, self.currentOrder + 1):
Avm = self.input.v('Av' + str(m), range(0, jmax + 1),
range(0, kmax + 1), range(0, fmax + 1))
if Avm is not None:
uz = self.input.d('u' + str(self.currentOrder - m),
range(0, jmax + 1),
range(0, kmax + 1),
range(0, fmax + 1),
dim='z')
ksi += ny.complexAmplitudeProduct(Avm, uz, 2)
ksi_z = ny.derivative(ksi, 'z', self.input.slice('grid'))
ksi = np.concatenate((np.zeros([jmax + 1, kmax + 1, fmax]), ksi),
2)
ksi_z = np.concatenate(
(np.zeros([jmax + 1, kmax + 1, fmax]), ksi_z), 2)
f_index += 1
f_names.append(['mixing', 'general'])
F[:, :, :, f_index] = ksi_z
Fsurf[:, :, :, f_index] = -ksi[:, [0], :]
if self.input.v('BottomBC') in ['PartialSlip']:
Fbed[:, :, :, f_index] = -ksi[:, [-1], :]
# 5.b higher order no-slip coefficient
for m in range(1, self.currentOrder + 1):
roughness = self.input.v('Roughness' + str(m),
range(0, jmax + 1), [0],
range(0, fmax + 1))
if roughness is not None:
u_str1 = 'u' + str(self.currentOrder - m)
u = self.input.v(u_str1, range(0, jmax + 1), [kmax],
range(0, fmax + 1))
ksi = ny.complexAmplitudeProduct(u, roughness, 2)
ksi = np.concatenate(
(np.zeros([jmax + 1, 1, fmax]), ksi), 2)
Fbed[:, :, :, f_index] = ksi
# 5.b Mixing-no-stress interaction
ksi = np.zeros([jmax + 1, 1, fmax + 1], dtype=complex)
for m in range(1, self.currentOrder + 1):
for k in range(0, self.currentOrder - m + 1):
for i in range(1, self.currentOrder - m - k + 1):
zetapermutations = self.multiindex(
m, self.currentOrder - m - k - i)
# a. make the zeta^m product
zetasum = 0
for perm in range(0, zetapermutations.shape[0]):
zetaname = 'zeta' + str(zetapermutations[perm, 0])
zeta = self.input.v(zetaname, range(0, jmax + 1),
[0], range(0, fmax + 1))
for comp in range(1, zetapermutations.shape[1]):
zetaname = 'zeta' + str(zetapermutations[perm,
comp])
zeta2 = self.input.v(zetaname,
range(0, jmax + 1), [0],
range(0, fmax + 1))
zeta = ny.complexAmplitudeProduct(
zeta, zeta2, 2)
zetasum = zetasum + zeta
# b. make the (Av*uz)^(m) term
Avuz = 0
for j in range(0, m + 1):
if j == 0:
Avder = self.input.v('Av' + str(i),
range(0, jmax + 1), [0],
range(0, fmax + 1))
else:
Avder = self.input.d('Av' + str(i),
range(0, jmax + 1), [0],
range(0, fmax + 1),
dim='z' * j)
if Avder is not None:
Avuz += scipy.misc.comb(
m, j) * ny.complexAmplitudeProduct(
Avder, self.surf_u_der[:, :, :, k,
m - j + 1],
2) # use m-j+1 as we need u_z^(m-j)
# c. add all to ksi
if not isinstance(Avuz, int):
ksi += 1. / np.math.factorial(
m) * ny.complexAmplitudeProduct(
Avuz, zetasum, 2)
ksi = np.concatenate((np.zeros([jmax + 1, 1, fmax]), ksi), 2)
f_index += 1
f_names.append(['mixing', 'no-stress'])
Fsurf[:, :, :, f_index] = -ksi
# 6. Stokes drift return flow
# determine RHS terms per submodule - next for water level
# Need to place this separate to make sure that stokes term are at the end.
# This way they will not be taken into account for velocity equation
if 'stokes' in self.submodulesToRun:
gamma = 0
for m in range(1, self.currentOrder + 1):
for k in range(0, self.currentOrder - m + 1):
zetapermutations = self.multiindex(
m, self.currentOrder - m - k)
# a. make the zeta^m product
zetasum = 0
for perm in range(0, zetapermutations.shape[0]):
zetaname = 'zeta' + str(zetapermutations[perm, 0])
zeta = self.input.v(zetaname, range(0, jmax + 1), [0],
range(0, fmax + 1))
for comp in range(1, zetapermutations.shape[1]):
zetaname = 'zeta' + str(zetapermutations[perm,
comp])
zeta2 = self.input.v(zetaname, range(0, jmax + 1),
[0], range(0, fmax + 1))
zeta = ny.complexAmplitudeProduct(zeta, zeta2, 2)
zetasum = zetasum + zeta
# b. make the (u)^(m-1) term
umz = self.surf_u_der[:, :, :, k, m - 1]
# c. add all to chi
gamma += 1. / np.math.factorial(
m) * ny.complexAmplitudeProduct(umz, zetasum, 2)
gamma = np.concatenate((np.zeros([jmax + 1, 1, fmax]), gamma), 2)
f_index += 1
f_names.append(['stokes', ''])
JuFirst[:, :, :, f_index] = gamma
## Solve equation
uCoef, uFirst[:, :, :, :
nRHSVelocity], uzCoef, uzFirst[:, :, :, :
nRHSVelocity], AMatrix = uFunctionMomentumConservative(
A,
F,
Fsurf,
Fbed,
self.input,
hasMatrix=
velocityMatrix)
if not velocityMatrix:
d['velocityMatrix'] = AMatrix
del AMatrix
################################################################################################################
# water level
################################################################################################################
## LHS terms
# try to get zetaMatrix from the input. If it is not found, proceed to calculate the matrix again
B = self.input.v(
'zetaMatrix'
) # known that this is numeric data, ask without arguments to retrieve full ndarray
zetaMatrix = True
if B is None:
zetaMatrix = False
utemp = uCoef.reshape(
uCoef.shape[:2] + (1, ) + uCoef.shape[2:]
) # reshape as the 'f' dimension is not grid conform; move it to a higher dimension
JuCoef = ny.integrate(utemp, 'z', kmax, 0,
self.input.slice('grid'))
JuCoef = JuCoef.reshape(jmax + 1, 1, ftot,
ftot) # reshape back to original grid
B = -G * JuCoef * self.input.v('B', np.arange(
0, jmax + 1)).reshape(jmax + 1, 1, 1, 1)
## RHS terms
# advection, no-stress, baroclinic
utemp = uFirst.reshape(
uFirst.shape[:2] + (1, ) + uFirst.shape[2:]
) # reshape as the 'f' dimension is not grid conform; move it to a higher dimension
JTemp = ny.integrate(utemp, 'z', kmax, 0, self.input.slice('grid'))
JuFirst += JTemp.reshape(
jmax + 1, 1, ftot,
uFirst.shape[-1]) # reshape back to original grid
BJuFirst = JuFirst * self.input.v('B', np.arange(0, jmax + 1)).reshape(
jmax + 1, 1, 1, 1)
# no open BC forcing & all terms in closed BC
Fopen = np.zeros([1, 1, ftot, nRHS], dtype=complex)
Fclosed = np.zeros([1, 1, ftot, nRHS], dtype=complex)
Fclosed += -JuFirst[jmax, 0, :, :] * self.input.v('B', jmax)
## Solve equation
zetaCoef, zetaxCoef, BMatrix = zetaFunctionMassConservative(
B, BJuFirst, Fopen, Fclosed, self.input, hasMatrix=zetaMatrix)
if not zetaMatrix:
d['zetaMatrix'] = BMatrix
del BMatrix
zetax = ny.eliminateNegativeFourier(zetaxCoef, 2)
zeta = ny.eliminateNegativeFourier(zetaCoef, 2)
################################################################################################################
# velocity
################################################################################################################
u = np.empty((jmax + 1, kmax + 1, ftot, nRHS), dtype=uCoef.dtype)
for j in range(0, jmax + 1):
u[j, :, :, :] = np.dot(uCoef[j, :, :, :],
-G * zetaxCoef[j, 0, :, :])
u += uFirst
u = ny.eliminateNegativeFourier(u, 2)
################################################################################################################
# Reduce number of components
################################################################################################################
# Select components by their velocity magnitude.
# This is measured as the 1-norm over z and f and the 2-norm over x
if maxContributions != 'all':
for submod in self.submodulesToRun:
f_names_numbers = zip(
f_names, range(0, len(f_names))
) # combine forcing names and its position in the 4th dimension of u and zeta
f_submod = [
f_names_numbers[i] for i in range(0, len(f_names))
if f_names[i][0] == submod
] # take only the forcing component of a particular submodule
# determine norm and sort the list
if f_submod:
unorm = [
np.linalg.norm(
np.linalg.norm(u[:, :, :, i], 1, (1, 2)), 2, 0)
for i in zip(*f_submod)[1]
]
sorted_fnames = [
list(l) for l in zip(*sorted(zip(unorm, f_submod),
key=lambda nrm: nrm[0]))
] # sorted list with entries (unorm, (forcing name, forcing position)) with smallest norm first
else:
unorm = []
sorted_fnames = []
# determine the contributions to be aggregated (redundant positions)
redundant_positions = [
sorted_fnames[1][i][1]
for i in range(0,
len(unorm) - maxContributions)
]
if len(redundant_positions) >= 1:
# replace first redundant element by sum and label 'other'
first_red_pos = redundant_positions[0]
u[:, :, :,
first_red_pos] = np.sum(u[:, :, :, redundant_positions],
3)
zeta[:, :, :, first_red_pos] = np.sum(
zeta[:, :, :, redundant_positions], 3)
f_names[first_red_pos][1] = 'other'
# remove other redundant positions
u = np.delete(u, redundant_positions[1:], 3)
zeta = np.delete(zeta, redundant_positions[1:], 3)
[
f_names.pop(i)
for i in sorted(redundant_positions[1:], reverse=True)
]
################################################################################################################
# vertical velocity
################################################################################################################
w = self.verticalVelocity(u)
################################################################################################################
# Make final dictionary to return
################################################################################################################
d['zeta' + str(self.currentOrder)] = {}
d['u' + str(self.currentOrder)] = {}
d['w' + str(self.currentOrder)] = {}
for submod in self.submodulesToRun:
if submod in zip(*f_names)[0]:
d['zeta' + str(self.currentOrder)][submod] = {}
d['u' + str(self.currentOrder)][submod] = {}
d['w' + str(self.currentOrder)][submod] = {}
for i, submod in enumerate(f_names):
nf = ny.functionTemplates.NumericalFunctionWrapper(
zeta[:, :, :, i], self.input.slice('grid'))
nf.addDerivative(zetax[:, :, :, i], 'x')
if submod == 'baroc':
d['zeta' + str(self.currentOrder)][submod[0]] = nf.function
d['u' + str(self.currentOrder)][submod[0]] = u[:, :, :, i]
d['w' + str(self.currentOrder)][submod[0]] = w[:, :, :, i]
else:
d['zeta' +
str(self.currentOrder)][submod[0]][submod[1]] = nf.function
d['u' +
str(self.currentOrder)][submod[0]][submod[1]] = u[:, :, :, i]
d['w' +
str(self.currentOrder)][submod[0]][submod[1]] = w[:, :, :, i]
d['surfder'] = self.surf_u_der
d['surfstress'] = self.surf_stress
return d
