def lu(a):
"""
Performe LU decomposition on the matrix a so A = L*U
"""
dw = True
if dw:
def loop_body(l, u):
c = get_iterator()
l[c:, c - 1] = (u[c:, c - 1] / u[c - 1, c - 1:c])
u[c:, c -
1:] = u[c:, c - 1:] - l[c:, c - 1][:, None] * u[c - 1, c - 1:]
u = a.copy()
l = np.zeros_like(a)
np.diagonal(l)[:] = 1.0
loop.do_while(loop_body, u.shape[0] - 2, l, u)
return (l, u)
else:
u = a.copy()
l = np.zeros_like(a)
np.diagonal(l)[:] = 1.0
for c in range(1, u.shape[0]):
l[c:, c - 1] = (u[c:, c - 1] / u[c - 1, c - 1:c])
u[c:, c -
1:] = u[c:, c - 1:] - l[c:, c - 1][:, None] * u[c - 1, c - 1:]
np.flush()
return (l, u)
def solve_tridiag(a, b, c, d, B):
assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape
#if not climate.is_bohrium:
# return lapack.dgtsv(a.flatten()[1:],b.flatten(),c.flatten()[:-1],d.flatten())[3].reshape(a.shape)
n = a.shape[-1]
x, cp, dp = np.zeros_like(a), np.zeros_like(a), np.zeros_like(a)
# initialize c-prime and d-prime
cp[..., 0] = c[..., 0] / b[..., 0]
dp[..., 0] = d[..., 0] / b[..., 0]
# solve for vectors c-prime and d-prime
def loop_body(cp, dp):
i = get_iterator(1)
m = b[..., i] - cp[..., i - 1] * a[..., i]
fxa = 1.0 / m
cp[..., i] = c[..., i] * fxa
dp[..., i] = (d[..., i] - dp[..., i - 1] * a[..., i]) * fxa
B.do_while(loop_body, n - 1, cp, dp)
x[..., n - 1] = dp[..., n - 1]
def loop_body(x):
i = get_iterator()
x[..., n - 2 -
i] = dp[..., n - 2 - i] - cp[..., n - 2 - i] * x[..., n - 1 - i]
B.do_while(loop_body, n - 1, x)
return x
def adv_flux_superbee_wgrid(adv_fe, adv_fn, adv_ft, var, u_wgrid, v_wgrid, w_wgrid, maskW, dxt, dyt, dzw, cost, cosu, dt_tracer):
"""
Calculates advection of a tracer defined on Wgrid
"""
maskUtr = bh.zeros_like(maskW)
maskUtr[:-1, :, :] = maskW[1:, :, :] * maskW[:-1, :, :]
adv_fe[...] = 0.
adv_fe[1:-2, 2:-2, :] = _adv_superbee(u_wgrid, var, maskUtr, dxt, 0, cost, cosu, dt_tracer)
maskVtr = bh.zeros_like(maskW)
maskVtr[:, :-1, :] = maskW[:, 1:, :] * maskW[:, :-1, :]
adv_fn[...] = 0.
adv_fn[2:-2, 1:-2, :] = _adv_superbee(v_wgrid, var, maskVtr, dyt, 1, cost, cosu, dt_tracer)
maskWtr = bh.zeros_like(maskW)
maskWtr[:, :, :-1] = maskW[:, :, 1:] * maskW[:, :, :-1]
adv_ft[...] = 0.
adv_ft[2:-2, 2:-2, :-1] = _adv_superbee(w_wgrid, var, maskWtr, dzw, 2, cost, cosu, dt_tracer)
def lu(a):
"""
Performe LU decomposition on the matrix a so A = L*U
"""
u = a.copy()
l = np.zeros_like(a)
np.diagonal(l)[:] = 1.0
for c in xrange(1,u.shape[0]):
l[c:,c-1] = (u[c:,c-1]/u[c-1,c-1:c])
u[c:,c-1:] = u[c:,c-1:] - l[c:,c-1][:,None] * u[c-1,c-1:]
np.flush(u)
return (l,u)
def lu(a):
"""
Performe LU decomposition on the matrix a so A = L*U
"""
u = a.copy()
l = np.zeros_like(a)
np.diagonal(l)[:] = 1.0
for c in xrange(1,u.shape[0]):
l[c:,c-1] = (u[c:,c-1]/u[c-1,c-1:c])
u[c:,c-1:] = u[c:,c-1:] - l[c:,c-1][:,None] * u[c-1,c-1:]
np.flush(u)
return (l,u)
def weighted_line(r0, c0, r1, c1, w, rmin=0, rmax=np.inf):
# The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0).
# If either of these cases are violated, do some switches.
if abs(c1-c0) < abs(r1-r0):
# Switch x and y, and switch again when returning.
xx, yy, val = weighted_line(c0, r0, c1, r1, w, rmin=rmin, rmax=rmax)
return (yy, xx, val)
# At this point we know that the distance in columns (x) is greater
# than that in rows (y). Possibly one more switch if c0 > c1.
if c0 > c1:
return weighted_line(r1, c1, r0, c0, w, rmin=rmin, rmax=rmax)
# The following is now always < 1 in abs
num=r1-r0
denom=c1-c0
slope = np.divide(num,denom,out=np.zeros_like(denom), where=denom!=0)
# Adjust weight by the slope
w *= np.sqrt(1+np.abs(slope)) / 2
# We write y as a function of x, because the slope is always <= 1
# (in absolute value)
x = np.arange(c0, c1+1, dtype=float)
y = x * slope + (c1*r0-c0*r1) / (c1-c0)
# Now instead of 2 values for y, we have 2*np.ceil(w/2).
# All values are 1 except the upmost and bottommost.
thickness = np.ceil(w/2)
yy = (np.floor(y).reshape(-1,1) + np.arange(-thickness-1,thickness+2).reshape(1,-1))
xx = np.repeat(x, yy.shape[1])
vals = trapez(yy, y.reshape(-1,1), w).flatten()
yy = yy.flatten()
# Exclude useless parts and those outside of the interval
# to avoid parts outside of the picture
mask = np.logical_and.reduce((yy >= rmin, yy < rmax, vals > 0))
return (yy[mask].astype(int), xx[mask].astype(int), vals[mask])
def integrate_tke(u, v, w, maskU, maskV, maskW, dxt, dxu, dyt, dyu, dzt, dzw, cost, cosu, kbot, kappaM, mxl, forc, forc_tke_surface, tke, dtke):
tau = 0
taup1 = 1
taum1 = 2
dt_tracer = 1
dt_mom = 1
AB_eps = 0.1
alpha_tke = 1.
c_eps = 0.7
K_h_tke = 2000.
flux_east = bh.zeros_like(maskU)
flux_north = bh.zeros_like(maskU)
flux_top = bh.zeros_like(maskU)
sqrttke = bh.sqrt(bh.maximum(0., tke[:, :, :, tau]))
"""
integrate Tke equation on W grid with surface flux boundary condition
"""
dt_tke = dt_mom # use momentum time step to prevent spurious oscillations
"""
vertical mixing and dissipation of TKE
"""
ks = kbot[2:-2, 2:-2] - 1
a_tri = bh.zeros_like(maskU[2:-2, 2:-2])
b_tri = bh.zeros_like(maskU[2:-2, 2:-2])
c_tri = bh.zeros_like(maskU[2:-2, 2:-2])
d_tri = bh.zeros_like(maskU[2:-2, 2:-2])
delta = bh.zeros_like(maskU[2:-2, 2:-2])
delta[:, :, :-1] = dt_tke / dzt[bh.newaxis, bh.newaxis, 1:] * alpha_tke * 0.5 \
* (kappaM[2:-2, 2:-2, :-1] + kappaM[2:-2, 2:-2, 1:])
a_tri[:, :, 1:-1] = -delta[:, :, :-2] / \
dzw[bh.newaxis, bh.newaxis, 1:-1]
a_tri[:, :, -1] = -delta[:, :, -2] / (0.5 * dzw[-1])
b_tri[:, :, 1:-1] = 1 + (delta[:, :, 1:-1] + delta[:, :, :-2]) / dzw[bh.newaxis, bh.newaxis, 1:-1] \
+ dt_tke * c_eps \
* sqrttke[2:-2, 2:-2, 1:-1] / mxl[2:-2, 2:-2, 1:-1]
b_tri[:, :, -1] = 1 + delta[:, :, -2] / (0.5 * dzw[-1]) \
+ dt_tke * c_eps / mxl[2:-2, 2:-
2, -1] * sqrttke[2:-2, 2:-2, -1]
b_tri_edge = 1 + delta / dzw[bh.newaxis, bh.newaxis, :] \
+ dt_tke * c_eps / mxl[2:-2, 2:-2, :] * sqrttke[2:-2, 2:-2, :]
c_tri[:, :, :-1] = -delta[:, :, :-1] / dzw[bh.newaxis, bh.newaxis, :-1]
d_tri[...] = tke[2:-2, 2:-2, :, tau] + dt_tke * forc[2:-2, 2:-2, :]
d_tri[:, :, -1] += dt_tke * forc_tke_surface[2:-2, 2:-2] / (0.5 * dzw[-1])
sol, water_mask = solve_implicit(ks, a_tri, b_tri, c_tri, d_tri, b_edge=b_tri_edge)
tke[2:-2, 2:-2, :, taup1] = where(water_mask, sol, tke[2:-2, 2:-2, :, taup1])
"""
Add TKE if surface density flux drains TKE in uppermost box
"""
tke_surf_corr = bh.zeros(maskU.shape[:2])
mask = tke[2:-2, 2:-2, -1, taup1] < 0.0
tke_surf_corr[2:-2, 2:-2] = where(
mask,
-tke[2:-2, 2:-2, -1, taup1] * 0.5 * dzw[-1] / dt_tke,
0.
)
tke[2:-2, 2:-2, -1, taup1] = bh.maximum(0., tke[2:-2, 2:-2, -1, taup1])
"""
add tendency due to lateral diffusion
"""
flux_east[:-1, :, :] = K_h_tke * (tke[1:, :, :, tau] - tke[:-1, :, :, tau]) \
/ (cost[bh.newaxis, :, bh.newaxis] * dxu[:-1, bh.newaxis, bh.newaxis]) * maskU[:-1, :, :]
flux_east[-1, :, :] = 0.
flux_north[:, :-1, :] = K_h_tke * (tke[:, 1:, :, tau] - tke[:, :-1, :, tau]) \
/ dyu[bh.newaxis, :-1, bh.newaxis] * maskV[:, :-1, :] * cosu[bh.newaxis, :-1, bh.newaxis]
flux_north[:, -1, :] = 0.
tke[2:-2, 2:-2, :, taup1] += dt_tke * maskW[2:-2, 2:-2, :] * \
((flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis])
+ (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis]))
"""
add tendency due to advection
"""
adv_flux_superbee_wgrid(
flux_east, flux_north, flux_top, tke[:, :, :, tau],
u[..., tau], v[..., tau], w[..., tau], maskW, dxt, dyt, dzw,
cost, cosu, dt_tracer
)
dtke[2:-2, 2:-2, :, tau] = maskW[2:-2, 2:-2, :] * (-(flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis])
- (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
/ (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis]))
dtke[:, :, 0, tau] += -flux_top[:, :, 0] / dzw[0]
dtke[:, :, 1:-1, tau] += - \
(flux_top[:, :, 1:-1] - flux_top[:, :, :-2]) / dzw[1:-1]
dtke[:, :, -1, tau] += - \
(flux_top[:, :, -1] - flux_top[:, :, -2]) / \
(0.5 * dzw[-1])
"""
Adam Bashforth time stepping
"""
tke[:, :, :, taup1] += dt_tracer * ((1.5 + AB_eps) * dtke[:, :, :, tau] - (0.5 + AB_eps) * dtke[:, :, :, taum1])
return tke, dtke, tke_surf_corr
def connected_components(a, connectivity=8):
b = np.array([connectivity], dtype=a.dtype)
c = np.zeros_like(a)
ufuncs.extmethod("opencv_connected_components", c, a, b)
return c
def isoneutral_diffusion_pre(maskT, maskU, maskV, maskW, dxt, dxu, dyt, dyu,
dzt, dzw, cost, cosu, salt, temp, zt, K_iso, K_11,
K_22, K_33, Ai_ez, Ai_nz, Ai_bx, Ai_by):
"""
Isopycnal diffusion for tracer
following functional formulation by Griffies et al
Code adopted from MOM2.1
"""
epsln = 1e-20
iso_slopec = 1e-3
iso_dslope = 1e-3
K_iso_steep = 50.
tau = 0
dTdx = bh.zeros_like(K_11)
dSdx = bh.zeros_like(K_11)
dTdy = bh.zeros_like(K_11)
dSdy = bh.zeros_like(K_11)
dTdz = bh.zeros_like(K_11)
dSdz = bh.zeros_like(K_11)
"""
drho_dt and drho_ds at centers of T cells
"""
drdT = maskT * get_drhodT(salt[:, :, :, tau], temp[:, :, :, tau],
bh.abs(zt))
drdS = maskT * get_drhodS(salt[:, :, :, tau], temp[:, :, :, tau],
bh.abs(zt))
"""
gradients at top face of T cells
"""
dTdz[:, :, :-1] = maskW[:, :, :-1] * \
(temp[:, :, 1:, tau] - temp[:, :, :-1, tau]) / \
dzw[bh.newaxis, bh.newaxis, :-1]
dSdz[:, :, :-1] = maskW[:, :, :-1] * \
(salt[:, :, 1:, tau] - salt[:, :, :-1, tau]) / \
dzw[bh.newaxis, bh.newaxis, :-1]
"""
gradients at eastern face of T cells
"""
dTdx[:-1, :, :] = maskU[:-1, :, :] * (temp[1:, :, :, tau] - temp[:-1, :, :, tau]) \
/ (dxu[:-1, bh.newaxis, bh.newaxis] * cost[bh.newaxis, :, bh.newaxis])
dSdx[:-1, :, :] = maskU[:-1, :, :] * (salt[1:, :, :, tau] - salt[:-1, :, :, tau]) \
/ (dxu[:-1, bh.newaxis, bh.newaxis] * cost[bh.newaxis, :, bh.newaxis])
"""
gradients at northern face of T cells
"""
dTdy[:, :-1, :] = maskV[:, :-1, :] * \
(temp[:, 1:, :, tau] - temp[:, :-1, :, tau]) \
/ dyu[bh.newaxis, :-1, bh.newaxis]
dSdy[:, :-1, :] = maskV[:, :-1, :] * \
(salt[:, 1:, :, tau] - salt[:, :-1, :, tau]) \
/ dyu[bh.newaxis, :-1, bh.newaxis]
def dm_taper(sx):
"""
tapering function for isopycnal slopes
"""
return 0.5 * (1. + bh.tanh((-bh.abs(sx) + iso_slopec) / iso_dslope))
"""
Compute Ai_ez and K11 on center of east face of T cell.
"""
diffloc = bh.zeros_like(K_11)
diffloc[1:-2, 2:-2,
1:] = 0.25 * (K_iso[1:-2, 2:-2, 1:] + K_iso[1:-2, 2:-2, :-1] +
K_iso[2:-1, 2:-2, 1:] + K_iso[2:-1, 2:-2, :-1])
diffloc[1:-2, 2:-2, 0] = 0.5 * \
(K_iso[1:-2, 2:-2, 0] + K_iso[2:-1, 2:-2, 0])
sumz = bh.zeros_like(K_11)[1:-2, 2:-2]
for kr in range(2):
ki = 0 if kr == 1 else 1
for ip in range(2):
drodxe = drdT[1 + ip:-2 + ip, 2:-2, ki:] * dTdx[1:-2, 2:-2, ki:] \
+ drdS[1 + ip:-2 + ip, 2:-2, ki:] * dSdx[1:-2, 2:-2, ki:]
drodze = drdT[1 + ip:-2 + ip, 2:-2, ki:] * dTdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None] \
+ drdS[1 + ip:-2 + ip, 2:-2, ki:] * \
dSdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None]
sxe = -drodxe / (bh.minimum(0., drodze) - epsln)
taper = dm_taper(sxe)
sumz[:, :, ki:] += dzw[bh.newaxis, bh.newaxis, :-1 + kr or None] * maskU[1:-2, 2:-2, ki:] \
* bh.maximum(K_iso_steep, diffloc[1:-2, 2:-2, ki:] * taper)
Ai_ez[1:-2, 2:-2, ki:, ip, kr] = taper * \
sxe * maskU[1:-2, 2:-2, ki:]
K_11[1:-2, 2:-2, :] = sumz / (4. * dzt[bh.newaxis, bh.newaxis, :])
"""
Compute Ai_nz and K_22 on center of north face of T cell.
"""
diffloc[...] = 0
diffloc[2:-2, 1:-2,
1:] = 0.25 * (K_iso[2:-2, 1:-2, 1:] + K_iso[2:-2, 1:-2, :-1] +
K_iso[2:-2, 2:-1, 1:] + K_iso[2:-2, 2:-1, :-1])
diffloc[2:-2, 1:-2, 0] = 0.5 * \
(K_iso[2:-2, 1:-2, 0] + K_iso[2:-2, 2:-1, 0])
sumz = bh.zeros_like(K_11)[2:-2, 1:-2]
for kr in range(2):
ki = 0 if kr == 1 else 1
for jp in range(2):
drodyn = drdT[2:-2, 1 + jp:-2 + jp, ki:] * dTdy[2:-2, 1:-2, ki:] + \
drdS[2:-2, 1 + jp:-2 + jp, ki:] * dSdy[2:-2, 1:-2, ki:]
drodzn = drdT[2:-2, 1 + jp:-2 + jp, ki:] * dTdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None] \
+ drdS[2:-2, 1 + jp:-2 + jp, ki:] * \
dSdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None]
syn = -drodyn / (bh.minimum(0., drodzn) - epsln)
taper = dm_taper(syn)
sumz[:, :, ki:] += dzw[bh.newaxis, bh.newaxis, :-1 + kr or None] \
* maskV[2:-2, 1:-2, ki:] * bh.maximum(K_iso_steep, diffloc[2:-2, 1:-2, ki:] * taper)
Ai_nz[2:-2, 1:-2, ki:, jp, kr] = taper * \
syn * maskV[2:-2, 1:-2, ki:]
K_22[2:-2, 1:-2, :] = sumz / (4. * dzt[bh.newaxis, bh.newaxis, :])
"""
compute Ai_bx, Ai_by and K33 on top face of T cell.
"""
sumx = bh.zeros_like(K_11)[2:-2, 2:-2, :-1]
sumy = bh.zeros_like(K_11)[2:-2, 2:-2, :-1]
for kr in range(2):
drodzb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdz[2:-2, 2:-2, :-1] \
+ drdS[2:-2, 2:-2, kr:-1 + kr or None] * dSdz[2:-2, 2:-2, :-1]
# eastward slopes at the top of T cells
for ip in range(2):
drodxb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None] \
+ drdS[2:-2, 2:-2, kr:-1 + kr or None] * \
dSdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None]
sxb = -drodxb / (bh.minimum(0., drodzb) - epsln)
taper = dm_taper(sxb)
sumx += dxu[1 + ip:-3 + ip, bh.newaxis, bh.newaxis] * \
K_iso[2:-2, 2:-2, :-1] * taper * \
sxb**2 * maskW[2:-2, 2:-2, :-1]
Ai_bx[2:-2, 2:-2, :-1, ip, kr] = taper * \
sxb * maskW[2:-2, 2:-2, :-1]
# northward slopes at the top of T cells
for jp in range(2):
facty = cosu[1 + jp:-3 + jp] * dyu[1 + jp:-3 + jp]
drodyb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None] \
+ drdS[2:-2, 2:-2, kr:-1 + kr or None] * \
dSdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None]
syb = -drodyb / (bh.minimum(0., drodzb) - epsln)
taper = dm_taper(syb)
sumy += facty[bh.newaxis, :, bh.newaxis] * K_iso[2:-2, 2:-2, :-1] \
* taper * syb**2 * maskW[2:-2, 2:-2, :-1]
Ai_by[2:-2, 2:-2, :-1, jp, kr] = taper * \
syb * maskW[2:-2, 2:-2, :-1]
K_33[2:-2, 2:-2, :-1] = sumx / (4 * dxt[2:-2, bh.newaxis, bh.newaxis]) + \
sumy / (4 * dyt[bh.newaxis, 2:-2, bh.newaxis]
* cost[bh.newaxis, 2:-2, bh.newaxis])
K_33[2:-2, 2:-2, -1] = 0.