def field_vert_int_bal(arr, dp):
"""Impose vertical balance to the field, i.e. column integral = 0.
Most frequently used for mass flux divergence to impose mass balance.
"""
pos = np.ma.where(arr > 0, arr, 0)
neg = np.ma.where(arr < 0, arr, 0)
pos_int = int_dp_g(pos, dp)[:, np.newaxis, :, :]
neg_int = int_dp_g(neg, dp)[:, np.newaxis, :, :]
arr.values = pos - (pos_int/neg_int)*neg
return arr
def field_vert_int_bal(arr, dp):
"""Impose vertical balance to the field, i.e. column integral = 0.
Most frequently used for mass flux divergence to impose mass balance.
"""
pos = np.ma.where(arr > 0, arr, 0)
neg = np.ma.where(arr < 0, arr, 0)
pos_int = int_dp_g(pos, dp)[:, np.newaxis, :, :]
neg_int = int_dp_g(neg, dp)[:, np.newaxis, :, :]
arr.values = pos - (pos_int / neg_int) * neg
return arr
def moisture_column_divg_with_adj2(q, ps, u, v, evap, precip, radius, dp,
freq='1M'):
"""Column flux divergence, with the field defined per unit mass of air."""
col_moist_divg = column_flux_divg_with_adj(q, ps, u, v, evap, precip,
radius, dp)
col_mass_divg = mass_column_divg_with_adj(ps, u, v, evap, precip, radius,
dp, freq=freq)
return col_moist_divg - col_mass_divg * int_dp_g(q, dp) / ps
def energy_column_divg_with_adj3(temp, z, q, q_ice, ps, u, v, evap, precip,
radius, dp, freq='1M'):
"""Column flux divergence, with the field defined per unit mass of air."""
en = energy(temp, z, q, q_ice, u, v)
col_energy_divg = column_flux_divg(en, u, v, radius, dp)
col_mass_divg = mass_column_divg_with_adj(ps, u, v, evap, precip, radius,
dp, freq=freq)
return col_energy_divg - col_mass_divg * int_dp_g(en, dp) / ps
def dry_mass_column_tendency(ps, q, dp, freq='1M'):
"""Combined time-tendency term in column mass budget equation.
See e.g. Trenberth 1991, Eq. 9.
"""
return (time_tendency_first_to_last(ps, freq=freq) -
grav.value * time_tendency_first_to_last(int_dp_g(q, dp),
freq=freq))
def energy_column_budget_residual(temp, z, q, q_ice, u, v, swdn_toa, swup_toa,
olr, swup_sfc, swdn_sfc, lwup_sfc, lwdn_sfc,
shflx, evap, dp, radius):
en = energy(temp, z, q, q_ice, u, v)
tendency = time_tendency_each_timestep(int_dp_g(en, dp))
transport = column_flux_divg(en, u, v, radius, dp)
source = energy_column_source(swdn_toa, swup_toa, olr, swup_sfc, swdn_sfc,
lwup_sfc, lwdn_sfc, shflx, evap)
return tendency + transport - source
def uv_mass_adjustment(uv, q, ps, dp):
"""Correction to subtract from outputted u or v to balance mass budget.
C.f. Trenberth 1991, J. Climate, equations 9 and 10.
"""
wvp = int_dp_g(q, dp)
numer = integrate((1. - q)*uv, dp, PFULL_STR)
denom = (ps - grav.value * wvp)
return numer / denom
def mse_budget_advec_residual(temp, hght, sphum, ucomp, vcomp, omega,
p, dp, radius, swdn_toa, swup_toa, olr, swup_sfc,
swdn_sfc, lwup_sfc, lwdn_sfc, shflx, evap):
"""Residual in vertically integrated MSE budget."""
mse_ = mse(temp, hght, sphum)
transport = field_total_advec(mse_, ucomp, vcomp, omega, p, radius)
trans_vert_int = int_dp_g(transport, dp)
f_net = column_energy(swdn_toa, swup_toa, olr, swup_sfc, swdn_sfc,
lwup_sfc, lwdn_sfc, shflx, evap)
return f_net - trans_vert_int
def moisture_column_budget_with_adj_lhs(ps, u, v, q, radius, dp, freq='1M'):
"""Column moisture budget LHS, with mass-balance adjustment applied.
The mass adjustment: subtract off the column mass residual (deduced from
the column dry mass budget) times the column average specific humidity.
"""
moist_lhs = moisture_column_budget_lhs(u, v, q, radius, dp, freq=freq)
dry_mass_resid = dry_mass_column_budget_residual(ps, u, v, q, radius, dp)
wvp = int_dp_g(q, dp)
return moist_lhs - dry_mass_resid * grav.value * wvp / ps
def energy_column_budget_adj_residual(temp, z, q, q_ice, u, v, swdn_toa,
swup_toa, olr, swup_sfc, swdn_sfc,
lwup_sfc, lwdn_sfc, shflx, evap, precip,
ps, dp, radius, freq='1M'):
en = energy(temp, z, q, q_ice, u, v)
tendency = time_tendency_each_timestep(int_dp_g(en, dp))
transport = energy_column_divg_adj(
temp, z, q, q_ice, u, v, swdn_toa, swup_toa, olr, swup_sfc, swdn_sfc,
lwup_sfc, lwdn_sfc, shflx, evap, precip, ps, dp, radius, freq='1M'
)
source = energy_column_source(swdn_toa, swup_toa, olr, swup_sfc, swdn_sfc,
lwup_sfc, lwdn_sfc, shflx, evap)
return tendency + transport - source
def energy_adj(temp, z, q, q_ice, u, v, swdn_toa, swup_toa, olr, swup_sfc,
swdn_sfc, lwup_sfc, lwdn_sfc, shflx, evap, dp, radius):
"""Adjustment to subtract off of u, v to impose column energy balance."""
residual = energy_column_budget_residual(
temp, z, q, q_ice, u, v, swdn_toa, swup_toa, olr, swup_sfc,
swdn_sfc, lwup_sfc, lwdn_sfc, shflx, evap, dp, radius
)
# Get the residual in terms of its spectral coefficients.
# Grab only one level of u, v since residual not defined vertically.
sph_int = SpharmInterface(u[:,0], v[:,0], rsphere=radius,
make_spharmt=True)
resid_spectral = SpharmInterface.prep_for_spharm(residual)
resid_spectral = sph_int.spharmt.grdtospec(resid_spectral)
# Assume residual stems entirely from divergent flow.
vort_spectral = np.zeros_like(resid_spectral)
u_adj, v_adj = sph_int.spharmt.getuv(vort_spectral, resid_spectral)
u_arr = sph_int.to_xray(u_adj)
v_arr = sph_int.to_xray(v_adj)
col_energy = int_dp_g(energy(temp, z, q, q_ice, u, v), dp)
return u_arr / col_energy, v_arr / col_energy
def d_dy_of_vert_int(arr, radius, dp):
return d_dy_from_lat(int_dp_g(arr, dp), radius, vec_field=True)
def d_dx_of_vert_int(arr, radius, dp):
return d_dx_from_latlon(int_dp_g(arr, dp), radius)
def energy_column(temp, z, q, q_ice, u, v, dp):
return int_dp_g(energy(temp, z, q, q_ice, u, v), dp)
def column_flux_divg(arr, u, v, radius, dp):
"""Column flux divergence, with the field defined per unit mass of air."""
return horiz_divg_spharm(int_dp_g(arr*u, dp), int_dp_g(arr*v, dp), radius)
def column_flux_divg_with_adj(arr, ps, u, v, evap, precip, radius, dp,
freq='1M'):
"""Column flux divergence, with the field defined per unit mass of air."""
u_adj, v_adj = mass_adjusted(ps, u, v, evap, precip, radius, dp, freq=freq)
return horiz_divg_spharm(int_dp_g(arr*u_adj, dp), int_dp_g(arr*v_adj, dp),
radius)
def moisture_column_tendency(q, dp, freq='1M'):
"""Time tendency of column-integrated water vapor."""
return time_tendency_first_to_last(int_dp_g(q, dp), freq=freq)
