def run_full_timestep_py(SOILTEMP, SOILMOIST, LWFLXNET_srfc, SWFLXNET_srfc,
SOILCP, SOILRHO, SOILDEPTH, OCEANMASK, TAIR_nz, QV_nz,
WIND_nz, RHO_nz, PSURF, COLP, WINDX_nz, WINDY_nz,
RAIN, DRAGCM, DRAGCH, A, dt):
# comute surface fluxes
if i_surface_fluxes:
SMOMXFLX, SMOMYFLX, SSHFLX, SLHFLX = calc_srfc_fluxes(
SOILTEMP, SOILMOIST, TAIR_nz, QV_nz, WIND_nz, RHO_nz, PSURF, COLP,
WINDX_nz, WINDY_nz, DRAGCM, DRAGCH, A, dt)
else:
SMOMXFLX = wp(0.)
SMOMYFLX = wp(0.)
SSHFLX = wp(0.)
SLHFLX = wp(0.)
# soil temperature change
if i_surface_SOILTEMP_tendency:
dSOILTEMPdt = tendency_SOILTEMP(LWFLXNET_srfc, SWFLXNET_srfc, SSHFLX,
SLHFLX, SOILCP, SOILRHO, SOILDEPTH)
else:
dSOILTEMPdt = wp(0.)
SOILTEMP = timestep_SOILTEMP(SOILTEMP, dSOILTEMPdt, dt)
# update surface albedo
SURFALBEDSW, SURFALBEDLW = calc_albedo(OCEANMASK, SOILTEMP, SOILMOIST)
# soil moisture change
SOILMOIST -= dt * SLHFLX / con_Lh
if SOILMOIST > max_moisture_soil:
SOILMOIST = max_moisture_soil
SOILMOIST += RAIN
return (SOILTEMP, SOILMOIST, SURFALBEDSW, SURFALBEDLW, SMOMXFLX, SMOMYFLX,
SSHFLX, SLHFLX)
def bulk_richardson_py(QV_k, POTT_k, POTT_km05, POTT_kp05, ALT_km05, ALT_kp05,
WINDX_km05, WINDX_kp05, WINDY_km05, WINDY_kp05):
POTT_v_k = calc_virtual_temperature(POTT_k, QV_k)
Ri_b_k = ((con_g / POTT_v_k * (POTT_km05 - POTT_kp05) *
(ALT_km05 - ALT_kp05)) / ((WINDX_km05 - WINDX_kp05)**wp(2.) +
(WINDY_km05 - WINDY_kp05)**wp(2.)))
return (Ri_b_k)
def initial_conditions(self, GR, **fields):
fields['OCEANMASK'][fields['HSURF'] > 100] = 0
fields['OCEANMASK'][fields['HSURF'] <= 100] = 1
fields['SOILDEPTH'][fields['OCEANMASK'] == 0] = depth_soil
fields['SOILDEPTH'][fields['OCEANMASK'] == 1] = depth_ocean
fields['SOILCP'][fields['OCEANMASK'] == 0] = cp_soil
fields['SOILCP'][fields['OCEANMASK'] == 1] = cp_ocean
fields['SOILRHO'][fields['OCEANMASK'] == 0] = rho_soil
fields['SOILRHO'][fields['OCEANMASK'] == 1] = rho_water
fields['SOILTEMP'][:, :,
0] = (295 - np.sin(GR.lat_rad[:, :, 0])**2 * 30)
fields['SOILTEMP'] -= fields['HSURF'] * wp(0.0100)
fields['SOILMOIST'][fields['OCEANMASK'] == 0] = moisture_soil
fields['SOILMOIST'][fields['OCEANMASK'] == 1] = moisture_ocean
fields['SURFALBEDSW'][fields['OCEANMASK'] == 1] = 0.05
fields['SURFALBEDLW'][fields['OCEANMASK'] == 1] = 0.00
fields['SURFALBEDSW'][fields['OCEANMASK'] == 0] = 0.20
fields['SURFALBEDLW'][fields['OCEANMASK'] == 0] = 0.00
fields['SURFALBEDSW'][fields['SOILTEMP'][:, :, 0] <= wp(273.15)] = 0.6
fields['SURFALBEDLW'][fields['SOILTEMP'][:, :, 0] <= wp(273.15)] = 0.0
return (fields)
def compute_K_coefs_py(HGT_k, WINDX_km05, WINDX_kp05, WINDY_km05, WINDY_kp05,
ALT_km05, ALT_kp05, QV_km05, QV_kp05, POTT_k, POTT_km05,
POTT_kp05):
# compute mixing length
mix_len = con_k * HGT_k / (wp(1.) + con_k * HGT_k / free_mix_len)
# bulk richardson number
QV_k = comp_VARVB_log(QV_kp05, QV_km05)
Ri_b_k = bulk_richardson(QV_k, POTT_k, POTT_km05, POTT_kp05, ALT_km05,
ALT_kp05, WINDX_km05, WINDX_kp05, WINDY_km05,
WINDY_kp05)
# vertical wind shear term
shear_term = sqrt(((WINDX_km05 - WINDX_kp05) /
(ALT_km05 - ALT_kp05))**wp(2.) +
((WINDY_km05 - WINDY_kp05) /
(ALT_km05 - ALT_kp05))**wp(2.))
# mixing coefficient for momentum
KMOM_k = mix_len**wp(2.) * shear_term * (Ri_c - Ri_b_k) / Ri_c
# negative values possible in above formula but implies zero.
# also most often the values are way too big ?? TODO
if KMOM_k < min_KMOM:
KMOM_k = min_KMOM
if KMOM_k > max_KMOM:
KMOM_k = max_KMOM
# mixing coefficient for heat and moisture
KHEAT_k = KMOM_k / con_Pr
## TODO
#KHEAT_k = wp(1.)
return (KMOM_k, KHEAT_k)
def calc_srfc_fluxes_py(SOILTEMP, SOILMOIST, TAIR_nz, QV_nz, WIND_nz, RHO_nz,
PSURF, COLP, WINDX_nz, WINDY_nz, DRAGCM, DRAGCH, A,
dt):
# surface momentum flux in x direction pointing towards atmosphere
SMOMXFLX = -DRAGCM * WIND_nz * WINDX_nz * COLP * A
# surface momentum flux in y direction pointing towards atmosphere
SMOMYFLX = -DRAGCM * WIND_nz * WINDY_nz * COLP * A
# surface sensible heat flux pointing towards atmosphere
# (w'theta')*rho*con_cp [W m-2]
SSHFLX = -DRAGCH * WIND_nz * (TAIR_nz - SOILTEMP) * RHO_nz * con_cp
# for QV of soil assume 60% of saturation specific humidity for SOILTEMP
QV_soil = calc_specific_humidity(SOILTEMP, wp(60.), PSURF)
# surface latent heat flux pointing towards atmosphere
# (w'qv')*rho*con_Lh [W m-2]
SLHFLX = -DRAGCH * WIND_nz * (QV_nz - QV_soil) * RHO_nz * con_Lh
# over land apply resistance to evaporation
if not math.isnan(SOILMOIST):
SLHFLX *= land_evap_resist
# only allow for positive latent heat flux (towards atmosphere)
SLHFLX = max(wp(0.), SLHFLX)
# reduce latent heat flux if no soil moisture is available over land
if not math.isnan(SOILMOIST):
if dt * SLHFLX / con_Lh > SOILMOIST:
SLHFLX = SOILMOIST * con_Lh / dt
return (SMOMXFLX, SMOMYFLX, SSHFLX, SLHFLX)
def interp_COLPA_js_py(COLP, COLP_jm1, COLP_im1, COLP_ip1, COLP_jm1_ip1,
COLP_jm1_im1, A, A_jm1, A_im1, A_ip1, A_jm1_ip1,
A_jm1_im1):
"""
"""
return (wp(1.) / wp(8.) * (COLP_jm1_ip1 * A_jm1_ip1 + COLP_ip1 * A_ip1 +
wp(2.) * COLP_jm1 * A_jm1 + wp(2.) * COLP * A +
COLP_jm1_im1 * A_jm1_im1 + COLP_im1 * A_im1))
def hor_adv_py(VAR, VAR_im1, VAR_ip1, VAR_jm1, VAR_jp1, UFLX, UFLX_ip1, VFLX,
VFLX_jp1, A):
"""
Horizontal advection using momentum fluxes UFLX and VFLX.
"""
return ((+UFLX * (VAR_im1 + VAR) / wp(2.) - UFLX_ip1 *
(VAR + VAR_ip1) / wp(2.) + VFLX *
(VAR_jm1 + VAR) / wp(2.) - VFLX_jp1 * (VAR + VAR_jp1) / wp(2.)) /
A)
def initial_conditions(self, GR, TAIR, PAIR, QV, RAINRATE, ACCRAIN):
for i in range(0, nx + 2 * nb):
for j in range(0, ny + 2 * nb):
for k in range(0, nz):
QV[i, j,
k] = calc_specific_humidity_py(TAIR[i, j, k], RH_init,
PAIR[i, j, k])
ACCRAIN[i, j, 0] = wp(0.)
RAINRATE[i, j, 0] = wp(0.)
GR.exchange_BC(QV)
def coriolis_and_spherical_VWIND_py(COLP, COLP_jm1, UWIND, UWIND_jm1,
UWIND_ip1, UWIND_ip1_jm1, corf, corf_jm1,
lat_rad, lat_rad_jm1, dlon_rad, dlat_rad):
"""
Calculate coriolis acceleration and spherical grid conversion
terms for VWIND.
"""
return (-con_rE * dlon_rad * dlat_rad / wp(2.) *
(COLP_jm1 * (UWIND_jm1 + UWIND_ip1_jm1) / wp(2.) *
(corf_jm1 * con_rE * cos(lat_rad_jm1) +
(UWIND_jm1 + UWIND_ip1_jm1) / wp(2.) * sin(lat_rad_jm1)) + COLP *
(UWIND + UWIND_ip1) / wp(2.) *
(corf * con_rE * cos(lat_rad) +
(UWIND + UWIND_ip1) / wp(2.) * sin(lat_rad))))
def calc_virtual_temperature_py(T, qv):
"""
Compute virtual temperature.
inputs:
T: temperature in K
qv: specific humidity in kg/kg
outputs:
T_v: virtual temperature in K
"""
epsilon = wp(0.622)
# TODO: do not assume mixing_ratio = qv
mixing_ratio = qv
T_v = T * (wp(1.) + mixing_ratio / epsilon) / (wp(1.) + mixing_ratio)
return (T_v)
def run_all_gpu(COLP, COLP_im1, COLP_ip1, COLP_jm1, COLP_jp1, COLP_im1_jm1,
COLP_im1_jp1, COLP_ip1_jm1, COLP_ip1_jp1, COLP_OLD,
COLP_OLD_im1, COLP_OLD_ip1, COLP_OLD_jm1, COLP_OLD_jp1,
COLP_OLD_im1_jm1, COLP_OLD_im1_jp1, COLP_OLD_ip1_jm1,
COLP_OLD_ip1_jp1, UWIND_OLD, dUFLXdt, VWIND_OLD, dVFLXdt,
POTT_OLD, dPOTTdt, QV_OLD, dQVdt, QC_OLD, dQCdt, A, A_im1,
A_ip1, A_jm1, A_jp1, A_im1_jm1, A_im1_jp1, A_ip1_jm1,
A_ip1_jp1, dt, i, j):
"""
"""
## UWIND
if j < ny + nb:
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)
UWIND = euler_forward_pw(UWIND_OLD, dUFLXdt, COLPA_is, COLPA_OLD_is,
dt)
## VWIND
if i < nx + nb:
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)
VWIND = euler_forward_pw(VWIND_OLD, dVFLXdt, COLPA_js, COLPA_OLD_js,
dt)
if i < nx + nb and j < ny + nb:
## POTT
POTT = euler_forward_pw(POTT_OLD, dPOTTdt, COLP, COLP_OLD, dt)
## QV
QV = euler_forward_pw(QV_OLD, dQVdt, COLP, COLP_OLD, dt)
# clip negative values
if QV < wp(0.):
QV = wp(0.)
## QC
QC = euler_forward_pw(QC_OLD, dQCdt, COLP, COLP_OLD, dt)
# clip negative values
if QC < wp(0.):
QC = wp(0.)
return (POTT, UWIND, VWIND, QV, QC)
def coriolis_and_spherical_UWIND_py(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):
"""
Calculate coriolis acceleration and spherical grid conversion
terms for UWIND.
"""
return (con_rE * dlon_rad * dlat_rad / wp(2.) *
(COLP_im1 * (VWIND_im1 + VWIND_im1_jp1) / wp(2.) *
(corf_is * con_rE * cos(lat_is_rad) +
(UWIND_im1 + UWIND) / wp(2.) * sin(lat_is_rad)) + COLP *
(VWIND + VWIND_jp1) / wp(2.) *
(corf_is * con_rE * cos(lat_is_rad) +
(UWIND + UWIND_ip1) / wp(2.) * sin(lat_is_rad))))
def pre_grad_py(PHI, PHI_dm1, COLP, COLP_dm1, POTT, POTT_dm1, PVTF, PVTF_dm1,
PVTFVB, PVTFVB_dm1, PVTFVB_dm1_kp1, PVTFVB_kp1, dsigma,
sigma_vb, sigma_vb_kp1, dgrid):
"""
Pressure gradient term for horizontal velocities.
dm1 & dp1 mean in the horizontal direction of interest
-1 & +1 grid point.
"""
return (-dgrid * ((PHI - PHI_dm1) * (COLP + COLP_dm1) / wp(2.) +
(COLP - COLP_dm1) * con_cp / wp(2.) *
(+POTT_dm1 / dsigma *
(sigma_vb_kp1 * (PVTFVB_dm1_kp1 - PVTF_dm1) + sigma_vb *
(PVTF_dm1 - PVTFVB_dm1)) + POTT / dsigma *
(sigma_vb_kp1 * (PVTFVB_kp1 - PVTF) + sigma_vb *
(PVTF - PVTFVB)))))
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 num_dif_pw_py(VAR, VAR_im1, VAR_ip1, VAR_jm1, VAR_jp1, COLP, COLP_im1,
COLP_ip1, COLP_jm1, COLP_jp1, VAR_dif_coef):
"""
Numerical diffusion with pressure weighting.
"""
return (VAR_dif_coef *
(+COLP_im1 * VAR_im1 + COLP_ip1 * VAR_ip1 + COLP_jm1 * VAR_jm1 +
COLP_jp1 * VAR_jp1 - wp(4.) * COLP * VAR))
def add_up_tendencies_py(POTT, POTT_im1, POTT_ip1, POTT_jm1, POTT_jp1,
POTT_km1, POTT_kp1, UFLX, UFLX_ip1, VFLX, VFLX_jp1,
COLP, COLP_im1, COLP_ip1, COLP_jm1, COLP_jp1, POTTVB,
POTTVB_kp1, WWIND, WWIND_kp1, COLP_NEW, PHI, PHI_kp1,
PHI_km1, PHIVB, PHIVB_kp1, KHEAT, KHEAT_kp1, RHO,
RHOVB, RHOVB_kp1, SSHFLX, dPOTTdt_RAD, A, dsigma,
POTT_dif_coef, k):
dPOTTdt = wp(0.)
if i_POTT_main_switch:
# HORIZONTAL ADVECTION
if i_POTT_hor_adv:
dPOTTdt = dPOTTdt + hor_adv(POTT, POTT_im1, POTT_ip1, POTT_jm1,
POTT_jp1, UFLX, UFLX_ip1, VFLX,
VFLX_jp1, A)
# VERTICAL ADVECTION
if i_POTT_vert_adv:
dPOTTdt = dPOTTdt + vert_adv(POTTVB, POTTVB_kp1, WWIND, WWIND_kp1,
COLP_NEW, dsigma, k)
## VERTICAL TURBULENT TRANSPORT
if i_POTT_vert_turb:
surf_flux = SSHFLX / con_cp
dPOTTdt_TURB = turb_flux_tendency(PHI, PHI_kp1, PHI_km1, PHIVB,
PHIVB_kp1, POTT, POTT_kp1,
POTT_km1, KHEAT, KHEAT_kp1, RHO,
RHOVB, RHOVB_kp1, COLP,
surf_flux, k)
dPOTTdt = dPOTTdt + dPOTTdt_TURB
# convert to [K hr-1] for user output
dPOTTdt_TURB = dPOTTdt_TURB / COLP * wp(3600.)
# NUMERICAL HORIZONTAL DIFUSION
if i_POTT_num_dif and (POTT_dif_coef > wp(0.)):
dPOTTdt = dPOTTdt + num_dif(POTT, POTT_im1, POTT_ip1, POTT_jm1,
POTT_jp1, COLP, COLP_im1, COLP_ip1,
COLP_jm1, COLP_jp1, POTT_dif_coef)
# RADIATION
if i_POTT_radiation:
dPOTTdt = dPOTTdt + radiation(dPOTTdt_RAD, COLP)
return (dPOTTdt, dPOTTdt_TURB)
def exchange_BC(self, FIELD):
"""
Python function (slow) to exchange boundaries.
Should not be used within time step loop but just for initializiation.
Advantage: Does not depend on import grid constants (e.g. nx)
Can therefore be used within grid.py and initial_conditiony.py
"""
dim2 = False
if len(FIELD.shape) == 2:
dim2 = True
fnx,fny = FIELD.shape
elif len(FIELD.shape) == 3:
fnx,fny,fnz = FIELD.shape
# zonal boundaries
if fnx == self.nxs+2*self.nb: # staggered in x
FIELD[0,::] = FIELD[self.nxs-1,::]
FIELD[self.nxs,::] = FIELD[1,::]
else: # unstaggered in x
FIELD[0,::] = FIELD[self.nx,::]
FIELD[self.nx+1,::] = FIELD[1,::]
if dim2:
# meridional boundaries
if fny == self.nys+2*self.nb: # staggered in y
for j in [0,1,self.nys,self.nys+1]:
FIELD[:,j] = wp(0.)
else: # unstaggered in y
FIELD[:,0] = FIELD[:,1]
FIELD[:,self.ny+1] = FIELD[:,self.ny]
else:
# meridional boundaries
if fny == self.nys+2*self.nb: # staggered in y
for j in [0,1,self.nys,self.nys+1]:
FIELD[:,j,:] = wp(0.)
else: # unstaggered in y
FIELD[:,0,:] = FIELD[:,1,:]
FIELD[:,self.ny+1,:] = FIELD[:,self.ny,:]
return(FIELD)
def tendency_SOILTEMP_py(LWFLXNET_srfc, SWFLXNET_srfc, SSHFLX, SLHFLX, SOILCP,
SOILRHO, SOILDEPTH):
dSOILTEMPdt = wp(0.)
if i_radiation:
dSOILTEMPdt += ((LWFLXNET_srfc + SWFLXNET_srfc) /
(SOILCP * SOILRHO * SOILDEPTH))
if i_surface_fluxes:
dSOILTEMPdt -= ((SSHFLX + SLHFLX) / (SOILCP * SOILRHO * SOILDEPTH))
return (dSOILTEMPdt)
def UVFLX_hor_adv_py(DWIND, DWIND_dm1, DWIND_dp1, DWIND_pm1, DWIND_pp1,
DWIND_dm1_pm1, DWIND_dm1_pp1, DWIND_dp1_pm1,
DWIND_dp1_pp1, BRFLX, BRFLX_dm1, CQFLX, CQFLX_pp1,
DSFLX_dm1, DSFLX_pp1, ETFLX, ETFLX_dm1_pp1,
sign_ETFLX_term):
"""
"""
return (+BRFLX_dm1 * (DWIND_dm1 + DWIND) / wp(2.) - BRFLX *
(DWIND + DWIND_dp1) / wp(2.) + CQFLX *
(DWIND_pm1 + DWIND) / wp(2.) - CQFLX_pp1 *
(DWIND + DWIND_pp1) / wp(2.) + DSFLX_dm1 *
(DWIND_dm1_pm1 + DWIND) / wp(2.) - DSFLX_pp1 *
(DWIND + DWIND_dp1_pp1) / wp(2.) + sign_ETFLX_term *
(+ETFLX * (DWIND_dp1_pm1 + DWIND) / wp(2.) - ETFLX_dm1_pp1 *
(DWIND + DWIND_dm1_pp1) / wp(2.)))
def interp_VAR_ds_py(VAR, VAR_dm1, VAR_pm1, VAR_pp1, VAR_pm1_dm1, VAR_pp1_dm1,
rigid_wall, p_ind, np):
"""
Interpolate VAR
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.
"""
# left rigid wall
if rigid_wall and (p_ind == nb):
VAR_ds = wp(0.25) * (VAR_pp1_dm1 + VAR_pp1 + VAR_dm1 + VAR) / con_g
# right rigid wall
elif rigid_wall and (p_ind == np):
VAR_ds = wp(0.25) * (VAR_dm1 + VAR + VAR_pm1_dm1 + VAR_pm1) / con_g
# inside domain (not at boundary of perpendicular dimension)
else:
VAR_ds = wp(0.125) * (VAR_pp1_dm1 + VAR_pp1 + wp(2.) * VAR_dm1 +
wp(2.) * VAR + VAR_pm1_dm1 + VAR_pm1) / con_g
return (VAR_ds)
def comp_VARVB_log_py(VAR, VAR_km1):
"""
Compute variable value at vertical border using
logarithmic interpolation.
(see Jacobson page 213.)
"""
min_val = wp(0.0000001)
VAR = max(VAR, min_val)
VAR_km1 = max(VAR_km1, min_val)
#VAR_kp1 = max(VAR_kp1, min_val)
if VAR_km1 == VAR:
VARVB = VAR
else:
VARVB = ((log(VAR_km1) - log(VAR)) / (wp(1.) / VAR - wp(1.) / VAR_km1))
#if VAR_kp1 == VAR:
# VARVB_kp1 = VAR
#else:
# VARVB_kp1 = ( ( log(VAR ) - log(VAR_kp1) ) /
# ( wp(1.) / VAR_kp1 - wp(1.) / VAR )
# )
return (VARVB)
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 interp_COLPA_is_py(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):
"""
"""
if j == 1:
return (wp(1.) / wp(4.) *
(COLP_im1_jp1 * A_im1_jp1 + COLP_jp1 * A_jp1 +
COLP_im1 * A_im1 + COLP * A))
elif j == ny:
return (wp(1.) / wp(4.) *
(COLP_im1_jm1 * A_im1_jm1 + COLP_jm1 * A_jm1 +
COLP_im1 * A_im1 + COLP * A))
else:
return (wp(1.) / wp(8.) *
(COLP_im1_jp1 * A_im1_jp1 + COLP_jp1 * A_jp1 +
wp(2.) * COLP_im1 * A_im1 + wp(2.) * COLP * A +
COLP_im1_jm1 * A_im1_jm1 + COLP_jm1 * A_jm1))
def calc_albedo_py(OCEANMASK, SOILTEMP, SOILMOIST):
# ocean
if OCEANMASK:
SURFALBEDSW = wp(0.05)
SURFALBEDLW = wp(0.00)
# sea ice
if SOILTEMP <= wp(273.15):
SURFALBEDSW = wp(0.5)
# land
else:
SURFALBEDLW = wp(0.0)
# desert
if SOILMOIST < desert_moisture_thresh:
SURFALBEDSW = wp(0.3)
else:
# snow / glacier
if SOILTEMP <= wp(273.15):
SURFALBEDSW = wp(0.6)
# normal land surface
else:
SURFALBEDSW = wp(0.2)
return (SURFALBEDSW, SURFALBEDLW)
def calc_momentum_fluxes_ij_py(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):
"""
"""
BFLX = wp(1.) / wp(12.) * (UFLX_jm1 + UFLX_ip1_jm1 + wp(2.) *
(UFLX + UFLX_ip1) + UFLX_jp1 + UFLX_ip1_jp1)
RFLX = wp(1.) / wp(12.) * (VFLX_im1 + VFLX_im1_jp1 + wp(2.) *
(VFLX + VFLX_jp1) + VFLX_ip1 + VFLX_ip1_jp1)
return (BFLX, RFLX)
def calc_momentum_fluxes_isj_py(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):
"""
"""
SFLX = wp(1.) / wp(24.) * (VFLX_im1 + VFLX_im1_jp1 + VFLX + VFLX_jp1 +
UFLX_im1 + wp(2.) * UFLX + UFLX_ip1)
TFLX = wp(1.) / wp(24.) * (VFLX_im1 + VFLX_im1_jp1 + VFLX + VFLX_jp1 -
UFLX_im1 - wp(2.) * UFLX - UFLX_ip1)
return (SFLX, TFLX)
def calc_momentum_fluxes_ijs_py(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):
"""
"""
DFLX = wp(1.) / wp(24.) * (VFLX_jm1 + wp(2.) * VFLX + VFLX_jp1 + UFLX_jm1 +
UFLX + UFLX_ip1_jm1 + UFLX_ip1)
EFLX = wp(1.) / wp(24.) * (VFLX_jm1 + wp(2.) * VFLX + VFLX_jp1 - UFLX_jm1 -
UFLX - UFLX_ip1_jm1 - UFLX_ip1)
return (DFLX, EFLX)
def exchange_BC_cpu(FIELD):
fnx, fny, fnz = FIELD.shape
# zonal boundaries
if fnx == nxs + 2 * nb: # staggered in x
FIELD[0, :, :] = FIELD[nxs - 1, :, :]
FIELD[nxs, :, :] = FIELD[1, :, :]
FIELD[nxs + 1, :, :] = FIELD[2, :, :]
else: # unstaggered in x
FIELD[0, :, :] = FIELD[nx, :, :]
FIELD[nx + 1, :, :] = FIELD[1, :, :]
# meridional boundaries
if fny == nys + 2 * nb: # staggered in y
for j in [0, 1, nys, nys + 1]:
FIELD[:, j, :] = wp(0.)
else: # unstaggered in y
FIELD[:, 0, :] = FIELD[:, 1, :]
FIELD[:, ny + 1, :] = FIELD[:, ny, :]
def calc_momentum_fluxes_isjs_py(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):
"""
"""
CFLX = wp(1.) / wp(12.) * (VFLX_im1_jm1 + VFLX_jm1 + wp(2.) *
(VFLX_im1 + VFLX) + VFLX_im1_jp1 + VFLX_jp1)
QFLX = wp(1.) / wp(12.) * (UFLX_im1_jm1 + UFLX_im1 + wp(2.) *
(UFLX_jm1 + UFLX) + UFLX_ip1_jm1 + UFLX_ip1)
return (CFLX, QFLX)
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)
