def interp_WWIND_UVWIND_py(DWIND, DWIND_km1, WWIND, WWIND_dm1, WWIND_pm1,
WWIND_pp1, WWIND_pm1_dm1, WWIND_pp1_dm1, COLP_NEW,
COLP_NEW_dm1, COLP_NEW_pm1, COLP_NEW_pp1,
COLP_NEW_pm1_dm1, COLP_NEW_pp1_dm1, A, A_dm1, A_pm1,
A_pp1, A_pm1_dm1, A_pp1_dm1, dsigma, dsigma_km1,
rigid_wall, p_ind, np, k):
"""
Interpolate WWIND * UVWIND (U or V, depending on direction)
onto position of repective horizontal wind.
d inds (e.g. dm1 = d minus 1) are in direction of hor. wind
vector.
p inds are perpendicular to direction of hor. wind vector.
if rigid_wall: Special BC for rigid wall parallel to flow
direction according to hint from Mark Jacobson during mail
conversation. np is number of grid cells and p_ind current
index in direction perpeindular to flow.
"""
if k == wp_int(0) or k == nzs - wp_int(1):
WWIND_DWIND = wp(0.)
else:
# left rigid wall
if rigid_wall and (p_ind == nb):
COLPAWWIND_ds_ks = wp(0.25) * (
COLP_NEW_pp1_dm1 * A_pp1_dm1 * WWIND_pp1_dm1 +
COLP_NEW_pp1 * A_pp1 * WWIND_pp1 +
COLP_NEW_dm1 * A_dm1 * WWIND_dm1 + COLP_NEW * A * WWIND)
# right rigid wall
elif rigid_wall and (p_ind == np):
COLPAWWIND_ds_ks = wp(0.25) * (
COLP_NEW_dm1 * A_dm1 * WWIND_dm1 + COLP_NEW * A * WWIND +
COLP_NEW_pm1_dm1 * A_pm1_dm1 * WWIND_pm1_dm1 +
COLP_NEW_pm1 * A_pm1 * WWIND_pm1)
# inside domain (not at boundary of perpendicular dimension)
else:
COLPAWWIND_ds_ks = wp(0.125) * (
COLP_NEW_pp1_dm1 * A_pp1_dm1 * WWIND_pp1_dm1 + COLP_NEW_pp1 *
A_pp1 * WWIND_pp1 + wp(2.) * COLP_NEW_dm1 * A_dm1 * WWIND_dm1 +
wp(2.) * COLP_NEW * A * WWIND + COLP_NEW_pm1_dm1 * A_pm1_dm1 *
WWIND_pm1_dm1 + COLP_NEW_pm1 * A_pm1 * WWIND_pm1)
# interpolate hor. wind on vertical interface
DWIND_ks = ((dsigma * DWIND_km1 + dsigma_km1 * DWIND) /
(dsigma + dsigma_km1))
# combine
WWIND_DWIND = COLPAWWIND_ds_ks * DWIND_ks
return (WWIND_DWIND)
def turb_flux_tendency_py(PHI, PHI_kp1, PHI_km1, PHIVB, PHIVB_kp1, VAR,
VAR_kp1, VAR_km1, KVAR, KVAR_kp1, RHO, RHOVB,
RHOVB_kp1, COLP, surf_flux_VAR, k):
"""
Vertical turbulent transport.
"""
ALT = PHI / con_g
ALT_kp1 = PHI_kp1 / con_g
ALT_km1 = PHI_km1 / con_g
ALTVB = PHIVB / con_g
ALTVB_kp1 = PHIVB_kp1 / con_g
if k == wp_int(0):
#dVARdt_TURB = wp(0.)
dVARdt_TURB = COLP * ((+wp(0.) -
((VAR - VAR_kp1) /
(ALT - ALT_kp1) * RHOVB_kp1 * KVAR_kp1)) /
((ALTVB - ALTVB_kp1) * RHO))
elif k == nz - 1:
#dVARdt_TURB = wp(0.)
dVARdt_TURB = COLP * (
(+((VAR_km1 - VAR) /
(ALT_km1 - ALT) * RHOVB * KVAR) + surf_flux_VAR) /
((ALTVB - ALTVB_kp1) * RHO))
else:
#dVARdt_TURB = wp(0.)
dVARdt_TURB = COLP * ((+((VAR_km1 - VAR) /
(ALT_km1 - ALT) * RHOVB * KVAR) -
((VAR - VAR_kp1) /
(ALT - ALT_kp1) * RHOVB_kp1 * KVAR_kp1)) /
((ALTVB - ALTVB_kp1) * RHO))
return (dVARdt_TURB)
def launch_numba_cpu(A, dsigma, POTT_dif_coef, dPOTTdt, POTT, UFLX, VFLX, COLP,
POTTVB, WWIND, COLP_NEW, PHI, PHIVB, KHEAT, RHO, RHOVB,
SSHFLX, dPOTTdt_TURB, dPOTTdt_RAD):
for i in prange(nb, nx + nb):
for j in range(nb, ny + nb):
for k in range(wp_int(0), nz):
dPOTTdt[i ,j ,k], dPOTTdt_TURB[i ,j ,k] = \
add_up_tendencies(POTT[i ,j ,k],
POTT [i-1,j ,k ], POTT [i+1,j ,k ],
POTT [i ,j-1,k ], POTT [i ,j+1,k ],
POTT [i ,j ,k-1], POTT [i ,j ,k+1],
UFLX [i ,j ,k ], UFLX[i+1,j ,k],
VFLX [i ,j ,k ], VFLX[i ,j+1,k],
COLP [i ,j ,0 ],
COLP [i-1,j ,0 ], COLP[i+1,j ,0],
COLP [i ,j-1,0 ], COLP[i ,j+1,0],
POTTVB [i ,j ,k ], POTTVB[i ,j ,k+1],
WWIND [i ,j ,k ], WWIND[i ,j ,k+1],
COLP_NEW [i ,j ,0 ],
PHI [i ,j ,k ], PHI [i ,j ,k+1],
PHI [i ,j ,k-1], PHIVB [i ,j ,k ],
PHIVB [i ,j ,k+1],
KHEAT [i ,j ,k ], KHEAT [i ,j ,k+1],
RHO [i ,j ,k ], RHOVB [i ,j ,k ],
RHOVB [i ,j ,k+1], SSHFLX [i ,j ,0 ],
dPOTTdt_RAD [i ,j ,k],
A [i ,j ,0 ],
dsigma [0 ,0 ,k ], POTT_dif_coef[0 ,0 ,k],
k)
def vert_adv_py(VARVB, VARVB_kp1, WWIND, WWIND_kp1, COLP_NEW, dsigma, k):
"""
Vertical advection.
"""
if k == wp_int(0):
return (COLP_NEW * (-WWIND_kp1 * VARVB_kp1) / dsigma)
elif k == nz:
return (COLP_NEW * (+WWIND * VARVB) / dsigma)
else:
return (COLP_NEW * (+WWIND * VARVB - WWIND_kp1 * VARVB_kp1) / dsigma)
def launch_numba_cpu(KMOM, KHEAT, PHIVB, HSURF, PHI, QV, WINDX, WINDY, POTTVB,
POTT):
for i in prange(0, nx + 2 * nb):
for j in range(0, ny + 2 * nb):
for k in range(wp_int(1), nzs - 1):
KMOM[i, j, k], KHEAT[i, j, k] = run_all(
PHIVB[i, j, k], HSURF[i, j, 0], PHI[i, j, k - 1],
PHI[i, j, k], QV[i, j, k - 1], QV[i, j, k], WINDX[i, j,
k - 1],
WINDX[i, j, k], WINDY[i, j, k - 1], WINDY[i, j, k],
POTTVB[i, j, k], POTT[i, j, k - 1], POTT[i, j, k], k)
def diag_secondary_cpu(POTTVB, TAIRVB, PVTFVB,
COLP, PAIR, PAIRVB, PHI, POTT,
TAIR, RHO, RHOVB, PVTF,
UWIND, VWIND, WINDX, WINDY, WIND):
for i in prange(0,nx+2*nb):
for j in range(0,ny+2*nb):
for k in range(wp_int(0),nzs):
PAIRVB[i,j,k] = wp(100000.)*(
PVTFVB[i,j,k])**(wp(1.)/con_kappa)
TAIRVB[i,j,k] = POTTVB[i,j,k] * PVTFVB[i,j,k]
RHOVB [i,j,k] = PAIRVB[i,j,k] / (con_Rd * TAIRVB[i,j,k])
for k in range(wp_int(0),nz):
TAIR [i,j,k] = POTT[i,j,k] * PVTF[i,j,k]
PAIR [i,j,k] = wp(100000.)*(PVTF[i,j,k])**(wp(1.)/con_kappa)
RHO [i,j,k] = PAIR[i,j,k] / (con_Rd * TAIR[i,j,k])
WINDX[i,j,k] = (UWIND[i ,j ,k] + UWIND[i+1,j ,k])/wp(2.)
WINDY[i,j,k] = (VWIND[i ,j ,k] + VWIND[i ,j+1,k])/wp(2.)
WIND [i,j,k] = (
((UWIND[i ,j ,k] + UWIND[i+1,j ,k])/wp(2.))**wp(2.) +
((VWIND[i ,j ,k] + VWIND[i ,j+1,k])/wp(2.))**wp(2.)
) ** (wp(1.)/wp(2.))
def launch_numba_cpu(UFLX, VFLX, FLXDIV, UWIND, VWIND, WWIND, COLP, dCOLPdt,
COLP_NEW, COLP_OLD, dyis, dxjs, dsigma, sigma_vb, A, dt):
for i in prange(nb, nx + nb):
for j in range(nb, ny + nb):
for k in range(wp_int(0), nz):
# MOMENTUM FLUXES
###############################################################
UFLX_i = calc_UFLX(UWIND[i, j, k], COLP[i, j, 0],
COLP[i - 1, j, 0], dyis[i, j, 0])
UFLX_ip1 = calc_UFLX(UWIND[i + 1, j, k], COLP[i + 1, j, 0],
COLP[i, j, 0], dyis[i + 1, j, 0])
VFLX_j = calc_VFLX(VWIND[i, j, k], COLP[i, j, 0],
COLP[i, j - 1, 0], dxjs[i, j, 0])
VFLX_jp1 = calc_VFLX(VWIND[i, j + 1, k], COLP[i, j + 1, 0],
COLP[i, j, 0], dxjs[i, j + 1, 0])
UFLX[i, j, k] = UFLX_i
VFLX[i, j, k] = VFLX_j
## MOMENTUM FLUX DIVERGENCE
###############################################################
FLXDIV[i, j, k] = calc_FLXDIV(UFLX_i, UFLX_ip1, VFLX_j,
VFLX_jp1, dsigma[0, 0,
k], A[i, j, 0])
## COLUMN PRESSURE TENDENCY
############################################################################
if i_COLP_main_switch:
dCOLPdt[:, :, 0] = -FLXDIV.sum(axis=2)
else:
dCOLPdt[:, :, 0] = wp(0.)
### PRESSURE TIME STEP
############################################################################
COLP_NEW[:, :, 0] = COLP_OLD[:, :, 0] + dt * dCOLPdt[:, :, 0]
### VERTICAL WIND
############################################################################
for i in prange(nb, nx + nb):
for j in range(nb, ny + nb):
flxdivsum = FLXDIV[i, j, 0]
for k in prange(1, nz):
WWIND[i, j, k] = (
-flxdivsum / COLP_NEW[i, j, 0] -
sigma_vb[0, 0, k] * dCOLPdt[i, j, 0] / COLP_NEW[i, j, 0])
flxdivsum += FLXDIV[i, j, k]
def diag_POTTVB_cpu(POTTVB, POTT, PVTF, PVTFVB):
for i in prange(0,nx+2*nb):
for j in range(0,ny+2*nb):
for k in range(wp_int(1),nzs-1):
POTTVB[i,j,k] = (
+ (PVTFVB[i,j,k] - PVTF [i,j,k-1]) * POTT[i,j,k-1]
+ (PVTF [i,j,k] - PVTFVB[i,j,k ]) * POTT[i,j,k ]
) / (PVTF[i,j,k] - PVTF[i,j,k-1])
if k == 1:
# extrapolate model top POTTVB
POTTVB[i,j,k-1] = POTT[i,j,k-1] - (
POTTVB[i,j,k] - POTT[i,j,k-1] )
elif k == nzs-2:
# extrapolate model bottom POTTVB
POTTVB[i,j,k+1] = POTT[i,j,k ] - (
POTTVB[i,j,k] - POTT[i,j,k ] )
def diag_PVTF_cpu(COLP, PVTF, PVTFVB, sigma_vb):
for i in prange(0,nx+2*nb):
for j in range(0,ny+2*nb):
for k in range(wp_int(0),nz):
pairvb_km12 = pair_top + sigma_vb[0,0,k ] * COLP[i,j,0]
pairvb_kp12 = pair_top + sigma_vb[0,0,k+1] * COLP[i,j,0]
PVTF[i,j,k] = wp(1.)/(wp(1.)+con_kappa) * (
pow( pairvb_kp12/wp(100000.) , con_kappa ) * pairvb_kp12 -
pow( pairvb_km12/wp(100000.) , con_kappa ) * pairvb_km12
) /( pairvb_kp12 - pairvb_km12 )
PVTFVB[i,j,k] = pow( pairvb_km12/wp(100000.) , con_kappa )
if k == nz-1:
PVTFVB[i,j,k+1] = pow( pairvb_kp12/wp(100000.) , con_kappa )
def launch_numba_cpu(A, dsigma, moist_dif_coef,
dQVdt, dQVdt_TURB, QV, dQCdt, QC, UFLX, VFLX, COLP,
WWIND, COLP_NEW,
PHI, PHIVB, KHEAT, RHO, RHOVB, SLHFLX):
for i in prange(nb,nx+nb):
for j in range(nb,ny+nb):
for k in range(wp_int(0),nz):
dQVdt[i ,j ,k], dQVdt_TURB[i ,j ,k], dQCdt[i ,j ,k] = \
add_up_tendencies(
QV [i ,j ,k ],
QV [i-1,j ,k ], QV [i+1,j ,k ],
QV [i ,j-1,k ], QV [i ,j+1,k ],
QV [i ,j ,k-1], QV [i ,j ,k+1],
QC [i ,j ,k ],
QC [i-1,j ,k ], QC [i+1,j ,k ],
QC [i ,j-1,k ], QC [i ,j+1,k ],
QC [i ,j ,k-1], QC [i ,j ,k+1],
UFLX [i ,j ,k ], UFLX [i+1,j ,k ],
VFLX [i ,j ,k ], VFLX [i ,j+1,k ],
COLP [i ,j ,0 ],
COLP [i-1,j ,0 ], COLP [i+1,j ,0 ],
COLP [i ,j-1,0 ], COLP [i ,j+1,0 ],
WWIND [i ,j ,k ], WWIND [i ,j ,k+1],
COLP_NEW [i ,j ,0 ],
PHI [i ,j ,k ], PHI [i ,j ,k+1],
PHI [i ,j ,k-1], PHIVB [i ,j ,k ],
PHIVB [i ,j ,k+1],
KHEAT [i ,j ,k ], KHEAT [i ,j ,k+1],
RHO [i ,j ,k ], RHOVB [i ,j ,k ],
RHOVB [i ,j ,k+1], SLHFLX [i ,j ,0 ],
A [i ,j ,0 ],
dsigma [0 ,0 ,k ], moist_dif_coef[0 ,0 ,k],
k)
def interp_KMOM_dUVWINDdz_py(
DWIND, DWIND_km1, KMOM, KMOM_dm1, KMOM_pm1, KMOM_pp1, KMOM_pm1_dm1,
KMOM_pp1_dm1, RHOVB, RHOVB_dm1, RHOVB_pm1, RHOVB_pp1, RHOVB_pm1_dm1,
RHOVB_pp1_dm1, PHI, PHI_dm1, PHI_pm1, PHI_pp1, PHI_pm1_dm1,
PHI_pp1_dm1, PHI_km1, PHI_km1_dm1, PHI_km1_pm1, PHI_km1_pp1,
PHI_km1_pm1_dm1, PHI_km1_pp1_dm1, COLP_NEW, COLP_NEW_dm1, COLP_NEW_pm1,
COLP_NEW_pp1, COLP_NEW_pm1_dm1, COLP_NEW_pp1_dm1, A, A_dm1, A_pm1,
A_pp1, A_pm1_dm1, A_pp1_dm1, dsigma, dsigma_km1, rigid_wall, p_ind, np,
k):
"""
Interpolate KMOM * UVWIND (U or V, depending on direction)
onto position of repective horizontal wind.
d inds (e.g. dm1 = d minus 1) are in direction of hor. wind
vector.
p inds are perpendicular to direction of hor. wind vector.
if rigid_wall: Special BC for rigid wall parallel to flow
direction according to hint from Mark Jacobson during mail
conversation. np is number of grid cells and p_ind current
index in direction perpeindular to flow.
"""
if k == wp_int(0) or k == nzs - wp_int(1):
KMOM_DWIND = wp(0.)
else:
# left rigid wall
if rigid_wall and (p_ind == nb):
COLPAKMOM_ds_ks = wp(0.25) * (
RHOVB_pp1_dm1 * COLP_NEW_pp1_dm1 * A_pp1_dm1 * KMOM_pp1_dm1 +
RHOVB_pp1 * COLP_NEW_pp1 * A_pp1 * KMOM_pp1 + RHOVB_dm1 *
COLP_NEW_dm1 * A_dm1 * KMOM_dm1 + RHOVB * COLP_NEW * A * KMOM)
ALT_ds_km1 = wp(0.25) * (PHI_km1_pp1_dm1 + PHI_km1_pp1 +
PHI_km1_dm1 + PHI_km1) / con_g
ALT_ds = wp(0.25) * (PHI_pp1_dm1 + PHI_pp1 + PHI_dm1 + PHI) / con_g
# right rigid wall
elif rigid_wall and (p_ind == np):
COLPAKMOM_ds_ks = wp(0.25) * (
RHOVB_dm1 * COLP_NEW_dm1 * A_dm1 * KMOM_dm1 +
RHOVB * COLP_NEW * A * KMOM +
RHOVB_pm1_dm1 * COLP_NEW_pm1_dm1 * A_pm1_dm1 * KMOM_pm1_dm1 +
RHOVB_pm1 * COLP_NEW_pm1 * A_pm1 * KMOM_pm1)
ALT_ds_km1 = wp(0.25) * (PHI_km1_dm1 + PHI_km1 + PHI_km1_pm1_dm1 +
PHI_km1_pm1) / con_g
ALT_ds = wp(0.25) * (PHI_dm1 + PHI + PHI_pm1_dm1 + PHI_pm1) / con_g
# inside domain (not at boundary of perpendicular dimension)
else:
COLPAKMOM_ds_ks = wp(0.125) * (
RHOVB_pp1_dm1 * COLP_NEW_pp1_dm1 * A_pp1_dm1 * KMOM_pp1_dm1 +
RHOVB_pp1 * COLP_NEW_pp1 * A_pp1 * KMOM_pp1 +
wp(2.) * RHOVB_dm1 * COLP_NEW_dm1 * A_dm1 * KMOM_dm1 +
wp(2.) * RHOVB * COLP_NEW * A * KMOM +
RHOVB_pm1_dm1 * COLP_NEW_pm1_dm1 * A_pm1_dm1 * KMOM_pm1_dm1 +
RHOVB_pm1 * COLP_NEW_pm1 * A_pm1 * KMOM_pm1)
ALT_ds_km1 = wp(0.125) * (PHI_km1_pp1_dm1 + PHI_km1_pp1 +
wp(2.) * PHI_km1_dm1 + wp(2.) * PHI_km1 +
PHI_km1_pm1_dm1 + PHI_km1_pm1) / con_g
ALT_ds = wp(0.125) * (PHI_pp1_dm1 + PHI_pp1 + wp(2.) * PHI_dm1 +
wp(2.) * PHI + PHI_pm1_dm1 + PHI_pm1) / con_g
# interpolate hor. wind on vertical interface
dDWINDdz_ks = ((DWIND_km1 - DWIND) / (ALT_ds_km1 - ALT_ds))
# combine
KMOM_dDWINDdz = COLPAKMOM_ds_ks * dDWINDdz_ks
return (KMOM_dDWINDdz)
def launch_numba_cpu_main(dUFLXdt, UFLX, UWIND, VWIND, BFLX_3D, CFLX_3D,
DFLX_3D, EFLX_3D, PHI, PHIVB, COLP, POTT, PVTF,
PVTFVB, WWIND_UWIND, KMOM_dUWINDdz, RHO,
dUFLXdt_TURB, SMOMXFLX, corf_is, lat_is_rad,
dlon_rad, dlat_rad, dyis, dsigma, sigma_vb,
UVFLX_dif_coef):
for i in prange(nb, nxs + nb):
for j in range(nb, ny + nb):
for k in range(wp_int(0), nz):
# Prepare momentum fluxes and set boundary conditions if
# necessary
BFLX = BFLX_3D[i, j, k]
CFLX = CFLX_3D[i, j, k]
EFLX = EFLX_3D[i, j, k]
DFLX_jp1 = DFLX_3D[i, j + 1, k]
CFLX_jp1 = CFLX_3D[i, j + 1, k]
# BCx i
if i == nb:
BFLX_im1 = BFLX_3D[nx, j, k]
DFLX_im1 = DFLX_3D[nx, j, k]
EFLX_im1_jp1 = EFLX_3D[nx, j + 1, k]
else:
BFLX_im1 = BFLX_3D[i - 1, j, k]
DFLX_im1 = DFLX_3D[i - 1, j, k]
EFLX_im1_jp1 = EFLX_3D[i - 1, j + 1, k]
# BCy js
if j == nb:
DFLX_im1 = wp(0.)
CFLX = wp(0.)
EFLX = wp(0.)
if j == ny + nb - 1:
DFLX_jp1 = wp(0.)
CFLX_jp1 = wp(0.)
EFLX_im1_jp1 = wp(0.)
dUFLXdt[i, j, k], dUFLXdt_TURB[i, j, k] = add_up_tendencies(
# 3D
UFLX[i, j, k],
UFLX[i - 1, j, k],
UFLX[i + 1, j, k],
UFLX[i, j - 1, k],
UFLX[i, j + 1, k],
VWIND[i, j, k],
VWIND[i - 1, j, k],
VWIND[i, j + 1, k],
VWIND[i - 1, j + 1, k],
UWIND[i, j, k],
UWIND[i - 1, j, k],
UWIND[i + 1, j, k],
UWIND[i, j - 1, k],
UWIND[i, j + 1, k],
UWIND[i - 1, j - 1, k],
UWIND[i - 1, j + 1, k],
UWIND[i + 1, j - 1, k],
UWIND[i + 1, j + 1, k],
BFLX,
BFLX_im1,
CFLX,
CFLX_jp1,
DFLX_im1,
DFLX_jp1,
EFLX,
EFLX_im1_jp1,
PHI[i, j, k],
PHI[i - 1, j, k],
POTT[i, j, k],
POTT[i - 1, j, k],
PVTF[i, j, k],
PVTF[i - 1, j, k],
PVTFVB[i, j, k],
PVTFVB[i - 1, j, k],
PVTFVB[i - 1, j, k + 1],
PVTFVB[i, j, k + 1],
WWIND_UWIND[i, j, k],
WWIND_UWIND[i, j, k + 1],
PHIVB[i, j, k],
PHIVB[i - 1, j, k],
PHIVB[i, j - 1, k],
PHIVB[i, j + 1, k],
PHIVB[i - 1, j - 1, k],
PHIVB[i - 1, j + 1, k],
PHIVB[i, j, k + 1],
PHIVB[i - 1, j, k + 1],
PHIVB[i, j - 1, k + 1],
PHIVB[i, j + 1, k + 1],
PHIVB[i - 1, j - 1, k + 1],
PHIVB[i - 1, j + 1, k + 1],
RHO[i, j, k],
RHO[i - 1, j, k],
RHO[i, j - 1, k],
RHO[i, j + 1, k],
RHO[i - 1, j - 1, k],
RHO[i - 1, j + 1, k],
SMOMXFLX[i, j, 0],
SMOMXFLX[i - 1, j, 0],
SMOMXFLX[i, j - 1, 0],
SMOMXFLX[i, j + 1, 0],
SMOMXFLX[i - 1, j - 1, 0],
SMOMXFLX[i - 1, j + 1, 0],
KMOM_dUWINDdz[i, j, k],
KMOM_dUWINDdz[i, j, k + 1],
# 2D
COLP[i, j, 0],
COLP[i - 1, j, 0],
# GR horizontal
corf_is[i, j, 0],
lat_is_rad[i, j, 0],
dlon_rad[i, j, 0],
dlat_rad[i, j, 0],
dyis[i, j, 0],
# GR vertical
dsigma[0, 0, k],
sigma_vb[0, 0, k],
sigma_vb[0, 0, k + 1],
UVFLX_dif_coef[0, 0, k],
k,
j)
def launch_numba_cpu_prep_adv(WWIND_UWIND, WWIND_VWIND,
UWIND, VWIND, WWIND,
UFLX_3D, VFLX_3D,
CFLX, QFLX, DFLX, EFLX,
SFLX, TFLX, BFLX, RFLX,
COLP, COLP_NEW,
KMOM_dUWINDdz, KMOM_dVWINDdz,
KMOM, PHI, RHOVB,
A, dsigma):
###########################################################################
if i_UVFLX_vert_adv:
for i in prange(nb,nxs+nb):
for j in range(nb,ny+nb):
for k in range(wp_int(0),nzs):
WWIND_UWIND[i ,j ,k ] = \
interp_WWIND_UVWIND(
UWIND [i ,j ,k ], UWIND [i ,j ,k-1],
WWIND [i ,j ,k ], WWIND [i-1,j ,k ],
WWIND [i ,j-1,k ], WWIND [i ,j+1,k ],
WWIND [i-1,j-1,k ], WWIND [i-1,j+1,k ],
COLP_NEW [i ,j ,0 ], COLP_NEW [i-1,j ,0 ],
COLP_NEW [i ,j-1,0 ], COLP_NEW [i ,j+1,0 ],
COLP_NEW [i-1,j-1,0 ], COLP_NEW [i-1,j+1,0 ],
A [i ,j ,0 ], A [i-1,j ,0 ],
A [i ,j-1,0 ], A [i ,j+1,0 ],
A [i-1,j-1,0 ], A [i-1,j+1,0 ],
dsigma [0 ,0 ,k ], dsigma [0 ,0 ,k-1],
True, j, ny, k)
for i in prange(nb,nx+nb):
for j in range(nb,nys+nb):
for k in range(wp_int(0),nzs):
WWIND_VWIND[i ,j ,k ] = \
interp_WWIND_UVWIND(
VWIND [i ,j ,k ], VWIND [i ,j ,k-1],
WWIND [i ,j ,k ], WWIND [i ,j-1,k ],
WWIND [i-1,j ,k ], WWIND [i+1,j ,k ],
WWIND [i-1,j-1,k ], WWIND [i+1,j-1,k ],
COLP_NEW [i ,j ,0 ], COLP_NEW [i ,j-1,0 ],
COLP_NEW [i-1,j ,0 ], COLP_NEW [i+1,j ,0 ],
COLP_NEW [i-1,j-1,0 ], COLP_NEW [i+1,j-1,0 ],
A [i ,j ,0 ], A [i ,j-1,0 ],
A [i-1,j ,0 ], A [i+1,j ,0 ],
A [i-1,j-1,0 ], A [i+1,j-1,0 ],
dsigma [0 ,0 ,k ], dsigma [0 ,0 ,k-1],
False, i, nx, k)
###########################################################################
if i_UVFLX_vert_turb:
for i in prange(nb,nxs+nb):
for j in range(nb,ny+nb):
for k in range(wp_int(0),nzs):
KMOM_dUWINDdz[i ,j ,k ] = \
interp_KMOM_dUVWINDdz(
UWIND [i ,j ,k ], UWIND [i ,j ,k-1],
KMOM [i ,j ,k ], KMOM [i-1,j ,k ],
KMOM [i ,j-1,k ], KMOM [i ,j+1,k ],
KMOM [i-1,j-1,k ], KMOM [i-1,j+1,k ],
RHOVB [i ,j ,k ], RHOVB [i-1,j ,k ],
RHOVB [i ,j-1,k ], RHOVB [i ,j+1,k ],
RHOVB [i-1,j-1,k ], RHOVB [i-1,j+1,k ],
PHI [i ,j ,k ], PHI [i-1,j ,k ],
PHI [i ,j-1,k ], PHI [i ,j+1,k ],
PHI [i-1,j-1,k ], PHI [i-1,j+1,k ],
PHI [i ,j ,k-1], PHI [i-1,j ,k-1],
PHI [i ,j-1,k-1], PHI [i ,j+1,k-1],
PHI [i-1,j-1,k-1], PHI [i-1,j+1,k-1],
COLP [i ,j ,0 ], COLP [i-1,j ,0 ],
COLP [i ,j-1,0 ], COLP [i ,j+1,0 ],
COLP [i-1,j-1,0 ], COLP [i-1,j+1,0 ],
A [i ,j ,0 ], A [i-1,j ,0 ],
A [i ,j-1,0 ], A [i ,j+1,0 ],
A [i-1,j-1,0 ], A [i-1,j+1,0 ],
dsigma [0 ,0 ,k ], dsigma [0 ,0 ,k-1],
True, j, ny, k)
for i in prange(nb,nx+nb):
for j in range(nb,nys+nb):
for k in range(wp_int(0),nzs):
KMOM_dVWINDdz[i ,j ,k ] = \
interp_KMOM_dUVWINDdz(
VWIND [i ,j ,k ], VWIND [i ,j ,k-1],
KMOM [i ,j ,k ], KMOM [i ,j-1,k ],
KMOM [i-1,j ,k ], KMOM [i+1,j ,k ],
KMOM [i-1,j-1,k ], KMOM [i+1,j-1,k ],
RHOVB [i ,j ,k ], RHOVB [i ,j-1,k ],
RHOVB [i-1,j ,k ], RHOVB [i+1,j ,k ],
RHOVB [i-1,j-1,k ], RHOVB [i+1,j-1,k ],
PHI [i ,j ,k ], PHI [i ,j-1,k ],
PHI [i-1,j ,k ], PHI [i+1,j ,k ],
PHI [i-1,j-1,k ], PHI [i+1,j-1,k ],
PHI [i ,j ,k-1], PHI [i ,j-1,k-1],
PHI [i-1,j ,k-1], PHI [i+1,j ,k-1],
PHI [i-1,j-1,k-1], PHI [i+1,j-1,k-1],
COLP [i ,j ,0 ], COLP [i ,j-1,0 ],
COLP [i-1,j ,0 ], COLP [i+1,j ,0 ],
COLP [i-1,j-1,0 ], COLP [i+1,j-1,0 ],
A [i ,j ,0 ], A [i ,j-1,0 ],
A [i-1,j ,0 ], A [i+1,j ,0 ],
A [i-1,j-1,0 ], A [i+1,j-1,0 ],
dsigma [0 ,0 ,k ], dsigma [0 ,0 ,k-1],
False, i, nx, k)
###########################################################################
if i_UVFLX_hor_adv:
for i in prange(nb,nxs+nb):
for j in range(nb,nys+nb):
for k in range(wp_int(0),nz):
UFLX = UFLX_3D[i ,j ,k ]
UFLX_im1 = UFLX_3D[i-1,j ,k ]
UFLX_im1_jm1 = UFLX_3D[i-1,j-1,k ]
UFLX_im1_jp1 = UFLX_3D[i-1,j+1,k ]
UFLX_ip1 = UFLX_3D[i+1,j ,k ]
UFLX_ip1_jm1 = UFLX_3D[i+1,j-1,k ]
UFLX_ip1_jp1 = UFLX_3D[i+1,j+1,k ]
UFLX_jm1 = UFLX_3D[i ,j-1,k ]
UFLX_jp1 = UFLX_3D[i ,j+1,k ]
VFLX = VFLX_3D[i ,j ,k ]
VFLX_im1 = VFLX_3D[i-1,j ,k ]
VFLX_im1_jm1 = VFLX_3D[i-1,j-1,k ]
VFLX_im1_jp1 = VFLX_3D[i-1,j+1,k ]
VFLX_ip1 = VFLX_3D[i+1,j ,k ]
VFLX_ip1_jm1 = VFLX_3D[i+1,j-1,k ]
VFLX_ip1_jp1 = VFLX_3D[i+1,j+1,k ]
VFLX_jm1 = VFLX_3D[i ,j-1,k ]
VFLX_jp1 = VFLX_3D[i ,j+1,k ]
if i < nxs+nb and j < nys+nb:
CFLX[i,j,k],QFLX[i,j,k] = \
calc_momentum_fluxes_isjs(
UFLX, UFLX_im1,
UFLX_im1_jm1, UFLX_im1_jp1,
UFLX_ip1, UFLX_ip1_jm1,
UFLX_ip1_jp1, UFLX_jm1,
UFLX_jp1,
VFLX, VFLX_im1,
VFLX_im1_jm1, VFLX_im1_jp1,
VFLX_ip1, VFLX_ip1_jm1,
VFLX_ip1_jp1, VFLX_jm1,
VFLX_jp1)
if i < nx+nb and j < nys+nb:
DFLX[i,j,k],EFLX[i,j,k] = \
calc_momentum_fluxes_ijs(
UFLX, UFLX_im1,
UFLX_im1_jm1, UFLX_im1_jp1,
UFLX_ip1, UFLX_ip1_jm1,
UFLX_ip1_jp1, UFLX_jm1,
UFLX_jp1,
VFLX, VFLX_im1,
VFLX_im1_jm1, VFLX_im1_jp1,
VFLX_ip1, VFLX_ip1_jm1,
VFLX_ip1_jp1, VFLX_jm1,
VFLX_jp1)
if i < nxs+nb and j < ny+nb:
SFLX[i,j,k],TFLX[i,j,k] = \
calc_momentum_fluxes_isj(
UFLX, UFLX_im1,
UFLX_im1_jm1, UFLX_im1_jp1,
UFLX_ip1, UFLX_ip1_jm1,
UFLX_ip1_jp1, UFLX_jm1,
UFLX_jp1,
VFLX, VFLX_im1,
VFLX_im1_jm1, VFLX_im1_jp1,
VFLX_ip1, VFLX_ip1_jm1,
VFLX_ip1_jp1, VFLX_jm1,
VFLX_jp1)
if i < nx+nb and j < ny+nb:
BFLX[i,j,k],RFLX[i,j,k] = \
calc_momentum_fluxes_ij(
UFLX, UFLX_im1,
UFLX_im1_jm1, UFLX_im1_jp1,
UFLX_ip1, UFLX_ip1_jm1,
UFLX_ip1_jp1, UFLX_jm1,
UFLX_jp1,
VFLX, VFLX_im1,
VFLX_im1_jm1, VFLX_im1_jp1,
VFLX_ip1, VFLX_ip1_jm1,
VFLX_ip1_jp1, VFLX_jm1,
VFLX_jp1)
def add_up_tendencies_py(
VFLX, VFLX_im1, VFLX_ip1, VFLX_jm1, VFLX_jp1, UWIND, UWIND_jm1,
UWIND_ip1, UWIND_ip1_jm1, VWIND, VWIND_jm1, VWIND_jp1, VWIND_im1,
VWIND_ip1, VWIND_jm1_im1, VWIND_jm1_ip1, VWIND_jp1_im1, VWIND_jp1_ip1,
RFLX, RFLX_jm1, QFLX, QFLX_ip1, SFLX_jm1, SFLX_ip1, TFLX, TFLX_jm1_ip1,
PHI, PHI_jm1, POTT, POTT_jm1, PVTF, PVTF_jm1, PVTFVB, PVTFVB_jm1,
PVTFVB_jm1_kp1, PVTFVB_kp1, WWIND_VWIND, WWIND_VWIND_kp1, PHIVB,
PHIVB_jm1, PHIVB_im1, PHIVB_ip1, PHIVB_im1_jm1, PHIVB_ip1_jm1,
PHIVB_kp1, PHIVB_kp1_jm1, PHIVB_kp1_im1, PHIVB_kp1_ip1,
PHIVB_kp1_im1_jm1, PHIVB_kp1_ip1_jm1, RHO, RHO_jm1, RHO_im1, RHO_ip1,
RHO_im1_jm1, RHO_ip1_jm1, SMOMYFLX, SMOMYFLX_jm1, SMOMYFLX_im1,
SMOMYFLX_ip1, SMOMYFLX_im1_jm1, SMOMYFLX_ip1_jm1, KMOM_dVWINDdz,
KMOM_dVWINDdz_kp1, COLP, COLP_jm1, corf, corf_jm1, lat_rad,
lat_rad_jm1, dlon_rad, dlat_rad, dxjs, dsigma, sigma_vb, sigma_vb_kp1,
UVFLX_dif_coef, k, i):
"""
Compute and add up all tendency contributions of VFLUX.
"""
dVFLXdt = wp(0.)
if i_UVFLX_main_switch:
# HORIZONTAL ADVECTION
if i_UVFLX_hor_adv:
#pass
dVFLXdt = dVFLXdt + UVFLX_hor_adv(
VWIND, VWIND_jm1, VWIND_jp1, VWIND_im1, VWIND_ip1,
VWIND_jm1_im1, VWIND_jm1_ip1, VWIND_jp1_im1, VWIND_jp1_ip1,
RFLX, RFLX_jm1, QFLX, QFLX_ip1, SFLX_jm1, SFLX_ip1, TFLX,
TFLX_jm1_ip1, wp(-1.))
# VERTICAL ADVECTION
if i_UVFLX_vert_adv:
dVFLXdt = dVFLXdt + ((WWIND_VWIND - WWIND_VWIND_kp1) / dsigma)
# VERTICAL TURBULENT TRANSPORT
if i_UVFLX_vert_turb:
ALTVB_js = interp_VAR_ds(PHIVB, PHIVB_jm1, PHIVB_im1, PHIVB_ip1,
PHIVB_im1_jm1, PHIVB_ip1_jm1, False, i,
nx) / con_g
ALTVB_kp1_js = interp_VAR_ds(
PHIVB_kp1, PHIVB_kp1_jm1, PHIVB_kp1_im1, PHIVB_kp1_ip1,
PHIVB_kp1_im1_jm1, PHIVB_kp1_ip1_jm1, False, i, nx) / con_g
RHO_js = interp_VAR_ds(RHO, RHO_jm1, RHO_im1, RHO_ip1, RHO_im1_jm1,
RHO_ip1_jm1, False, i, nx)
SMOMYFLX_js = interp_VAR_ds(SMOMYFLX, SMOMYFLX_jm1, SMOMYFLX_im1,
SMOMYFLX_ip1, SMOMYFLX_im1_jm1,
SMOMYFLX_ip1_jm1, False, i, nx)
if k == wp_int(0):
dVFLXdt_TURB = ((wp(0.) - KMOM_dVWINDdz_kp1) /
((ALTVB_js - ALTVB_kp1_js) * RHO_js))
elif k == wp_int(nz - 1):
dVFLXdt_TURB = ((KMOM_dVWINDdz + SMOMYFLX_js) /
((ALTVB_js - ALTVB_kp1_js) * RHO_js))
#dVFLXdt_TURB = wp(0.)
else:
dVFLXdt_TURB = ((KMOM_dVWINDdz - KMOM_dVWINDdz_kp1) /
((ALTVB_js - ALTVB_kp1_js) * RHO_js))
#dVFLXdt_TURB = wp(0.)
dVFLXdt = dVFLXdt + dVFLXdt_TURB
# CORIOLIS AND SPHERICAL GRID CONVERSION
if i_UVFLX_coriolis:
dVFLXdt = dVFLXdt + coriolis_and_spherical_VWIND(
COLP, COLP_jm1, UWIND, UWIND_jm1, UWIND_ip1, UWIND_ip1_jm1,
corf, corf_jm1, lat_rad, lat_rad_jm1, dlon_rad, dlat_rad)
# PRESSURE GRADIENT
if i_UVFLX_pre_grad:
dVFLXdt = dVFLXdt + pre_grad(
PHI, PHI_jm1, COLP, COLP_jm1, POTT, POTT_jm1, PVTF, PVTF_jm1,
PVTFVB, PVTFVB_jm1, PVTFVB_jm1_kp1, PVTFVB_kp1, dsigma,
sigma_vb, sigma_vb_kp1, dxjs)
# NUMERICAL HORIZONTAL DIFUSION
if i_UVFLX_num_dif and (UVFLX_dif_coef > wp(0.)):
dVFLXdt = dVFLXdt + num_dif(VFLX, VFLX_im1, VFLX_ip1, VFLX_jm1,
VFLX_jp1, UVFLX_dif_coef)
return (dVFLXdt, dVFLXdt_TURB)
def launch_numba_cpu_main(dVFLXdt, VFLX, UWIND, VWIND, RFLX_3D, SFLX_3D,
TFLX_3D, QFLX_3D, PHI, PHIVB, COLP, POTT, PVTF,
PVTFVB, WWIND_VWIND, KMOM_dVWINDdz, RHO,
dVFLXdt_TURB, SMOMYFLX, corf, lat_rad, dlon_rad,
dlat_rad, dxjs, dsigma, sigma_vb, UVFLX_dif_coef):
for i in prange(nb, nx + nb):
for j in range(nb, nys + nb):
for k in range(wp_int(0), nz):
# Prepare momentum fluxes and set boundary conditions if
# necessary
RFLX = RFLX_3D[i, j, k]
QFLX = QFLX_3D[i, j, k]
TFLX = TFLX_3D[i, j, k]
RFLX_jm1 = RFLX_3D[i, j - 1, k]
SFLX_jm1 = SFLX_3D[i, j - 1, k]
# BCx is
if i == nxs - 1:
QFLX_ip1 = QFLX_3D[1, j, k]
TFLX_ip1_jm1 = TFLX_3D[1, j - 1, k]
SFLX_ip1 = SFLX_3D[1, j, k]
else:
QFLX_ip1 = QFLX_3D[i + 1, j, k]
TFLX_ip1_jm1 = TFLX_3D[i + 1, j - 1, k]
SFLX_ip1 = SFLX_3D[i + 1, j, k]
#dVFLXdt[i ,j ,k] = add_up_tendencies(
dVFLXdt[i, j, k], dVFLXdt_TURB[i, j, k] = add_up_tendencies(
# 3D
VFLX[i, j, k],
VFLX[i - 1, j, k],
VFLX[i + 1, j, k],
VFLX[i, j - 1, k],
VFLX[i, j + 1, k],
UWIND[i, j, k],
UWIND[i, j - 1, k],
UWIND[i + 1, j, k],
UWIND[i + 1, j - 1, k],
VWIND[i, j, k],
VWIND[i, j - 1, k],
VWIND[i, j + 1, k],
VWIND[i - 1, j, k],
VWIND[i + 1, j, k],
VWIND[i - 1, j - 1, k],
VWIND[i + 1, j - 1, k],
VWIND[i - 1, j + 1, k],
VWIND[i + 1, j + 1, k],
RFLX,
RFLX_jm1,
QFLX,
QFLX_ip1,
SFLX_jm1,
SFLX_ip1,
TFLX,
TFLX_ip1_jm1,
PHI[i, j, k],
PHI[i, j - 1, k],
POTT[i, j, k],
POTT[i, j - 1, k],
PVTF[i, j, k],
PVTF[i, j - 1, k],
PVTFVB[i, j, k],
PVTFVB[i, j - 1, k],
PVTFVB[i, j - 1, k + 1],
PVTFVB[i, j, k + 1],
WWIND_VWIND[i, j, k],
WWIND_VWIND[i, j, k + 1],
PHIVB[i, j, k],
PHIVB[i, j - 1, k],
PHIVB[i - 1, j, k],
PHIVB[i + 1, j, k],
PHIVB[i - 1, j - 1, k],
PHIVB[i + 1, j - 1, k],
PHIVB[i, j, k + 1],
PHIVB[i, j - 1, k + 1],
PHIVB[i - 1, j, k + 1],
PHIVB[i + 1, j, k + 1],
PHIVB[i - 1, j - 1, k + 1],
PHIVB[i + 1, j - 1, k + 1],
RHO[i, j, k],
RHO[i, j - 1, k],
RHO[i - 1, j, k],
RHO[i + 1, j, k],
RHO[i - 1, j - 1, k],
RHO[i + 1, j - 1, k],
SMOMYFLX[i, j, 0],
SMOMYFLX[i, j - 1, 0],
SMOMYFLX[i - 1, j, 0],
SMOMYFLX[i + 1, j, 0],
SMOMYFLX[i - 1, j - 1, 0],
SMOMYFLX[i + 1, j - 1, 0],
KMOM_dVWINDdz[i, j, k],
KMOM_dVWINDdz[i, j, k + 1],
# 2D
COLP[i, j, 0],
COLP[i, j - 1, 0],
# GR horizontal
corf[i, j, 0],
corf[i, j - 1, 0],
lat_rad[i, j, 0],
lat_rad[i, j - 1, 0],
dlon_rad[i, j, 0],
dlat_rad[i, j, 0],
dxjs[i, j, 0],
# GR vertical
dsigma[0, 0, k],
sigma_vb[0, 0, k],
sigma_vb[0, 0, k + 1],
UVFLX_dif_coef[0, 0, k],
k,
i)
moist_dif_coef)
from io_constants import con_rE, con_omega
from io_restart import load_restart_grid
from io_initial_conditions import set_up_sigma_levels
from misc_utilities import Timer
###############################################################################
###############################################################################
# GLOBAL SETTINGS
###############################################################################
# vertical extent of cuda shared arrays
# (must be called before setting nz = wp_int(nz))
shared_nz = nz
# domain extent to be imported in modules
nz = wp_int(nz)
nzs = wp_int(nz+1)
nx = wp_int((lon1_deg - lon0_deg)/dlon_deg)
nxs = wp_int(nx+1)
ny = wp_int((lat1_deg - lat0_deg)/dlat_deg)
nys = wp_int(ny+1)
nb = wp_int(nb)
###############################################################################
# GPU COMPUTATION SETTINGS
###############################################################################
if gpu_enable:
if nz > 32:
raise NotImplementedError('More than 32 vertical levels '+
'not yet implemented for GPU.')
def make_timestep_cpu(COLP, COLP_OLD, UWIND, UWIND_OLD, dUFLXdt, VWIND,
VWIND_OLD, dVFLXdt, POTT, POTT_OLD, dPOTTdt, QV, QV_OLD,
dQVdt, QC, QC_OLD, dQCdt, A, dt):
for i in prange(nb, nxs + nb):
for j in range(nb, nys + nb):
COLP_ = COLP[i, j, 0]
COLP_im1 = COLP[i - 1, j, 0]
COLP_ip1 = COLP[i + 1, j, 0]
COLP_jm1 = COLP[i, j - 1, 0]
COLP_jp1 = COLP[i, j + 1, 0]
COLP_im1_jm1 = COLP[i - 1, j - 1, 0]
COLP_im1_jp1 = COLP[i - 1, j + 1, 0]
COLP_ip1_jm1 = COLP[i + 1, j - 1, 0]
COLP_ip1_jp1 = COLP[i + 1, j + 1, 0]
COLP_OLD_ = COLP_OLD[i, j, 0]
COLP_OLD_im1 = COLP_OLD[i - 1, j, 0]
COLP_OLD_ip1 = COLP_OLD[i + 1, j, 0]
COLP_OLD_jm1 = COLP_OLD[i, j - 1, 0]
COLP_OLD_jp1 = COLP_OLD[i, j + 1, 0]
COLP_OLD_im1_jm1 = COLP_OLD[i - 1, j - 1, 0]
COLP_OLD_im1_jp1 = COLP_OLD[i - 1, j + 1, 0]
COLP_OLD_ip1_jm1 = COLP_OLD[i + 1, j - 1, 0]
COLP_OLD_ip1_jp1 = COLP_OLD[i + 1, j + 1, 0]
A_ = A[i, j, 0]
A_im1 = A[i - 1, j, 0]
A_ip1 = A[i + 1, j, 0]
A_jm1 = A[i, j - 1, 0]
A_jp1 = A[i, j + 1, 0]
A_im1_jm1 = A[i - 1, j - 1, 0]
A_im1_jp1 = A[i - 1, j + 1, 0]
A_ip1_jm1 = A[i + 1, j - 1, 0]
A_ip1_jp1 = A[i + 1, j + 1, 0]
## UWIND
COLPA_is = interp_COLPA_is(COLP_, COLP_im1, COLP_jm1, COLP_jp1,
COLP_im1_jp1, COLP_im1_jm1, A_, A_im1,
A_jm1, A_jp1, A_im1_jp1, A_im1_jm1, j)
COLPA_OLD_is = interp_COLPA_is(COLP_OLD_, COLP_OLD_im1,
COLP_OLD_jm1, COLP_OLD_jp1,
COLP_OLD_im1_jp1, COLP_OLD_im1_jm1,
A_, A_im1, A_jm1, A_jp1, A_im1_jp1,
A_im1_jm1, j)
# VWIND
COLPA_js = interp_COLPA_js(COLP_, COLP_jm1, COLP_im1, COLP_ip1,
COLP_ip1_jm1, COLP_im1_jm1, A_, A_jm1,
A_im1, A_ip1, A_ip1_jm1, A_im1_jm1)
COLPA_OLD_js = interp_COLPA_js(COLP_OLD_, COLP_OLD_jm1,
COLP_OLD_im1, COLP_OLD_ip1,
COLP_OLD_ip1_jm1, COLP_OLD_im1_jm1,
A_, A_jm1, A_im1, A_ip1, A_ip1_jm1,
A_im1_jm1)
for k in range(wp_int(0), nz):
UWIND_OLD_ = UWIND_OLD[i, j, k]
dUFLXdt_ = dUFLXdt[i, j, k]
VWIND_OLD_ = VWIND_OLD[i, j, k]
dVFLXdt_ = dVFLXdt[i, j, k]
POTT_OLD_ = POTT_OLD[i, j, k]
dPOTTdt_ = dPOTTdt[i, j, k]
QV_OLD_ = QV_OLD[i, j, k]
dQVdt_ = dQVdt[i, j, k]
QC_OLD_ = QC_OLD[i, j, k]
dQCdt_ = dQCdt[i, j, k]
UWIND[i, j, k] = euler_forward_pw(UWIND_OLD_, dUFLXdt_,
COLPA_is, COLPA_OLD_is, dt)
VWIND[i, j, k] = euler_forward_pw(VWIND_OLD_, dVFLXdt_,
COLPA_js, COLPA_OLD_js, dt)
POTT[i, j, k] = euler_forward_pw(POTT_OLD_, dPOTTdt_, COLP_,
COLP_OLD_, dt)
QV[i, j, k] = euler_forward_pw(QV_OLD_, dQVdt_, COLP_,
COLP_OLD_, dt)
QC[i, j, k] = euler_forward_pw(QC_OLD_, dQCdt_, COLP_,
COLP_OLD_, dt)
def add_up_tendencies_py(
UFLX, UFLX_im1, UFLX_ip1, UFLX_jm1, UFLX_jp1, VWIND, VWIND_im1,
VWIND_jp1, VWIND_im1_jp1, UWIND, UWIND_im1, UWIND_ip1, UWIND_jm1,
UWIND_jp1, UWIND_im1_jm1, UWIND_im1_jp1, UWIND_ip1_jm1, UWIND_ip1_jp1,
BFLX, BFLX_im1, CFLX, CFLX_jp1, DFLX_im1, DFLX_jp1, EFLX, EFLX_im1_jp1,
PHI, PHI_im1, POTT, POTT_im1, PVTF, PVTF_im1, PVTFVB, PVTFVB_im1,
PVTFVB_im1_kp1, PVTFVB_kp1, WWIND_UWIND, WWIND_UWIND_kp1, PHIVB,
PHIVB_im1, PHIVB_jm1, PHIVB_jp1, PHIVB_jm1_im1, PHIVB_jp1_im1,
PHIVB_kp1, PHIVB_kp1_im1, PHIVB_kp1_jm1, PHIVB_kp1_jp1,
PHIVB_kp1_jm1_im1, PHIVB_kp1_jp1_im1, RHO, RHO_im1, RHO_jm1, RHO_jp1,
RHO_jm1_im1, RHO_jp1_im1, SMOMXFLX, SMOMXFLX_im1, SMOMXFLX_jm1,
SMOMXFLX_jp1, SMOMXFLX_jm1_im1, SMOMXFLX_jp1_im1, KMOM_dUWINDdz,
KMOM_dUWINDdz_kp1, COLP, COLP_im1, corf_is, lat_is_rad, dlon_rad,
dlat_rad, dyis, dsigma, sigma_vb, sigma_vb_kp1, UVFLX_dif_coef, k, j):
"""
Compute and add up all tendency contributions of UFLUX.
"""
dUFLXdt = wp(0.)
if i_UVFLX_main_switch:
# HORIZONTAL ADVECTION
if i_UVFLX_hor_adv:
dUFLXdt = dUFLXdt + UVFLX_hor_adv(
UWIND, UWIND_im1, UWIND_ip1, UWIND_jm1, UWIND_jp1,
UWIND_im1_jm1, UWIND_im1_jp1, UWIND_ip1_jm1, UWIND_ip1_jp1,
BFLX, BFLX_im1, CFLX, CFLX_jp1, DFLX_im1, DFLX_jp1, EFLX,
EFLX_im1_jp1, wp(1.))
# VERTICAL ADVECTION
if i_UVFLX_vert_adv:
dUFLXdt = dUFLXdt + ((WWIND_UWIND - WWIND_UWIND_kp1) / dsigma)
# VERTICAL TURBULENT TRANSPORT
if i_UVFLX_vert_turb:
ALTVB_is = interp_VAR_ds(PHIVB, PHIVB_im1, PHIVB_jm1, PHIVB_jp1,
PHIVB_jm1_im1, PHIVB_jp1_im1, True, j,
ny) / con_g
ALTVB_kp1_is = interp_VAR_ds(
PHIVB_kp1, PHIVB_kp1_im1, PHIVB_kp1_jm1, PHIVB_kp1_jp1,
PHIVB_kp1_jm1_im1, PHIVB_kp1_jp1_im1, True, j, ny) / con_g
RHO_is = interp_VAR_ds(RHO, RHO_im1, RHO_jm1, RHO_jp1, RHO_jm1_im1,
RHO_jp1_im1, True, j, ny)
SMOMXFLX_is = interp_VAR_ds(SMOMXFLX, SMOMXFLX_im1, SMOMXFLX_jm1,
SMOMXFLX_jp1, SMOMXFLX_jm1_im1,
SMOMXFLX_jp1_im1, True, j, ny)
if k == wp_int(0):
dUFLXdt_TURB = ((wp(0.) - KMOM_dUWINDdz_kp1) /
((ALTVB_is - ALTVB_kp1_is) * RHO_is))
elif k == wp_int(nz - 1):
dUFLXdt_TURB = ((KMOM_dUWINDdz + SMOMXFLX_is) /
((ALTVB_is - ALTVB_kp1_is) * RHO_is))
#dUFLXdt_TURB = wp(0.)
else:
dUFLXdt_TURB = ((KMOM_dUWINDdz - KMOM_dUWINDdz_kp1) /
((ALTVB_is - ALTVB_kp1_is) * RHO_is))
#dUFLXdt_TURB = wp(0.)
dUFLXdt = dUFLXdt + dUFLXdt_TURB
# CORIOLIS AND SPHERICAL GRID CONVERSION
if i_UVFLX_coriolis:
dUFLXdt = dUFLXdt + coriolis_and_spherical_UWIND(
COLP, COLP_im1, VWIND, VWIND_im1, VWIND_jp1, VWIND_im1_jp1,
UWIND, UWIND_im1, UWIND_ip1, corf_is, lat_is_rad, dlon_rad,
dlat_rad)
# PRESSURE GRADIENT
if i_UVFLX_pre_grad:
dUFLXdt = dUFLXdt + pre_grad(
PHI, PHI_im1, COLP, COLP_im1, POTT, POTT_im1, PVTF, PVTF_im1,
PVTFVB, PVTFVB_im1, PVTFVB_im1_kp1, PVTFVB_kp1, dsigma,
sigma_vb, sigma_vb_kp1, dyis)
# NUMERICAL HORIZONTAL DIFUSION
if i_UVFLX_num_dif and (UVFLX_dif_coef > wp(0.)):
dUFLXdt = dUFLXdt + num_dif(UFLX, UFLX_im1, UFLX_ip1, UFLX_jm1,
UFLX_jp1, UVFLX_dif_coef)
return (dUFLXdt, dUFLXdt_TURB)