def do_dyn(state, comm, timer=NullTimer()):
grid = spec.grid
copy_stencil(
state.delp,
state.dp1,
origin=grid.default_origin(),
domain=grid.domain_shape_standard(),
)
print("DynCore", grid.rank)
with timer.clock("DynCore"):
dyn_core.compute(state, comm)
if not spec.namelist.inline_q and state.nq != 0:
if spec.namelist.z_tracer:
print("Tracer2D1L", grid.rank)
with timer.clock("TracerAdvection"):
tracer_2d_1l.compute(
comm,
state.tracers,
state.dp1,
state.mfxd,
state.mfyd,
state.cxd,
state.cyd,
state.mdt,
state.nq,
)
else:
raise Exception("tracer_2d no =t implemented, turn on z_tracer")
def compute(qin, qout, kstart=0, nk=None, replace=False):
if nk is None:
nk = grid().npz - kstart
extrapolate_corners(qin, qout, kstart, nk)
if spec.namelist.grid_type < 3:
compute_qout_edges(qin, qout, kstart, nk)
qx = compute_qx(qin, qout, kstart, nk)
qy = compute_qy(qin, qout, kstart, nk)
qxx = compute_qxx(qx, qout, kstart, nk)
qyy = compute_qyy(qy, qout, kstart, nk)
compute_qout(qxx, qyy, qout, kstart, nk)
if replace:
copy_stencil(
qout,
qin,
origin=(grid().is_, grid().js, kstart),
domain=(grid().ie - grid().is_ + 2, grid().je - grid().js + 2,
nk),
)
else:
raise Exception("grid_type >= 3 is not implemented")
def compute(state, nq, dt):
tracers_dict(
state) # TODO get rid of this when finalize representation of tracers
grid = spec.grid
rdt = 1.0 / dt
k_bot = spec.namelist.n_sponge
if k_bot is not None:
if k_bot < 3:
return
else:
k_bot = grid.npz
if k_bot < min(grid.npz, 24):
t_max = T2_MAX
else:
t_max = T3_MAX
if state.pe[grid.is_, grid.js, 0] < 2.0:
t_min = T1_MIN
else:
t_min = T2_MIN
if spec.namelist.nwat == 0:
xvir = 0.0
# rz = 0 # hydrostatic only
else:
xvir = ZVIR
# rz = constants.RV_GAS - constants.RDGAS # hydrostatic only
m = 3
fra = dt / float(spec.namelist.fv_sg_adj)
if spec.namelist.hydrostatic:
raise Exception("Hydrostatic not supported for fv_subgridz")
q0 = {}
for tracername in utils.tracer_variables:
q0[tracername] = copy(state.__dict__[tracername])
origin = grid.compute_origin()
shape = state.delp.shape
u0 = utils.make_storage_from_shape(shape, origin)
v0 = utils.make_storage_from_shape(shape, origin)
w0 = utils.make_storage_from_shape(shape, origin)
gzh = utils.make_storage_from_shape(shape, origin)
gz = utils.make_storage_from_shape(shape, origin)
t0 = utils.make_storage_from_shape(shape, origin)
pm = utils.make_storage_from_shape(shape, origin)
hd = utils.make_storage_from_shape(shape, origin)
te = utils.make_storage_from_shape(shape, origin)
den = utils.make_storage_from_shape(shape, origin)
qcon = utils.make_storage_from_shape(shape, origin)
cvm = utils.make_storage_from_shape(shape, origin)
cpm = utils.make_storage_from_shape(shape, origin)
kbot_domain = (grid.nic, grid.njc, k_bot)
origin = grid.compute_origin()
init(
den,
gz,
gzh,
t0,
pm,
u0,
v0,
w0,
hd,
cvm,
cpm,
te,
state.ua,
state.va,
state.w,
state.pt,
state.peln,
state.delp,
state.delz,
q0["qvapor"],
q0["qliquid"],
q0["qrain"],
q0["qice"],
q0["qsnow"],
q0["qgraupel"],
xvir,
origin=origin,
domain=kbot_domain,
)
ri = utils.make_storage_from_shape(shape, origin)
ri_ref = utils.make_storage_from_shape(shape, origin)
mc = utils.make_storage_from_shape(shape, origin)
h0 = utils.make_storage_from_shape(shape, origin)
pt1 = utils.make_storage_from_shape(shape, origin)
pt2 = utils.make_storage_from_shape(shape, origin)
tv2 = utils.make_storage_from_shape(shape, origin)
ratios = {0: 0.25, 1: 0.5, 2: 0.999}
for n in range(m):
ratio = ratios[n]
compute_qcon(
qcon,
q0["qliquid"],
q0["qrain"],
q0["qice"],
q0["qsnow"],
q0["qgraupel"],
origin=origin,
domain=kbot_domain,
)
for k in range(k_bot - 1, 0, -1):
korigin = (grid.is_, grid.js, k)
korigin_m1 = (grid.is_, grid.js, k - 1)
kdomain = (grid.nic, grid.njc, 1)
kdomain_m1 = (grid.nic, grid.njc, 2)
m_loop(
ri,
ri_ref,
pm,
u0,
v0,
w0,
t0,
hd,
gz,
qcon,
state.delp,
state.pkz,
q0["qvapor"],
pt1,
pt2,
tv2,
t_min,
t_max,
ratio,
xvir,
origin=korigin,
domain=kdomain,
)
if k == 1:
ri_ref *= 4.0
if k == 2:
ri_ref *= 2.0
if k == 3:
ri_ref *= 1.5
# work around that gt4py will not accept interval(3, 4), no longer
# used, mc calc per k.
# m_loop_hack_interval_3_4(ri, ri_ref, mc, state.delp, ratio,
# origin=(grid.is_, grid.js, 1), domain=(grid.nic, grid.njc, k_bot - 1))
equivalent_mass_flux(ri,
ri_ref,
mc,
state.delp,
ratio,
origin=korigin,
domain=kdomain)
for tracername in utils.tracer_variables:
KH_instability_adjustment(
ri,
ri_ref,
mc,
q0[tracername],
state.delp,
h0,
origin=korigin,
domain=kdomain,
)
recompute_qcon(
ri,
ri_ref,
qcon,
q0["qliquid"],
q0["qrain"],
q0["qice"],
q0["qsnow"],
q0["qgraupel"],
origin=korigin_m1,
domain=kdomain,
)
KH_instability_adjustment(ri,
ri_ref,
mc,
u0,
state.delp,
h0,
origin=korigin,
domain=kdomain)
KH_instability_adjustment(ri,
ri_ref,
mc,
v0,
state.delp,
h0,
origin=korigin,
domain=kdomain)
KH_instability_adjustment(ri,
ri_ref,
mc,
w0,
state.delp,
h0,
origin=korigin,
domain=kdomain)
KH_instability_adjustment_te(ri,
ri_ref,
mc,
te,
state.delp,
h0,
hd,
origin=korigin,
domain=kdomain)
double_adjust_cvm(
cvm,
cpm,
gz,
u0,
v0,
w0,
hd,
t0,
te,
q0["qliquid"],
q0["qvapor"],
q0["qice"],
q0["qsnow"],
q0["qrain"],
q0["qgraupel"],
origin=korigin_m1,
domain=kdomain_m1,
)
if fra < 1.0:
fraction_adjust(
t0,
state.pt,
u0,
state.ua,
v0,
state.va,
w0,
state.w,
fra,
spec.namelist.hydrostatic,
origin=origin,
domain=kbot_domain,
)
for tracername in utils.tracer_variables:
fraction_adjust_tracer(
q0[tracername],
state.tracers[tracername],
fra,
origin=origin,
domain=kbot_domain,
)
for tracername in utils.tracer_variables:
copy_stencil(q0[tracername],
state.tracers[tracername],
origin=origin,
domain=kbot_domain)
finalize(
u0,
v0,
w0,
t0,
state.ua,
state.va,
state.pt,
state.w,
state.u_dt,
state.v_dt,
rdt,
origin=origin,
domain=kbot_domain,
)
def compute(comm, tracers, dp1, mfxd, mfyd, cxd, cyd, mdt, nq):
grid = spec.grid
shape = mfxd.data.shape
# start HALO update on q (in dyn_core in fortran -- just has started when
# this function is called...)
xfx = utils.make_storage_from_shape(shape, origin=grid.compute_x_origin())
yfx = utils.make_storage_from_shape(shape, origin=grid.compute_y_origin())
fx = utils.make_storage_from_shape(shape, origin=grid.compute_origin())
fy = utils.make_storage_from_shape(shape, origin=grid.compute_origin())
ra_x = utils.make_storage_from_shape(shape, origin=grid.compute_x_origin())
ra_y = utils.make_storage_from_shape(shape, origin=grid.compute_y_origin())
cmax = utils.make_storage_from_shape(shape, origin=grid.compute_origin())
dp2 = utils.make_storage_from_shape(shape, origin=grid.compute_origin())
flux_x(
cxd,
grid.dxa,
grid.dy,
grid.sin_sg3,
grid.sin_sg1,
xfx,
origin=grid.compute_x_origin(),
domain=grid.domain_y_compute_xbuffer(),
)
flux_y(
cyd,
grid.dya,
grid.dx,
grid.sin_sg4,
grid.sin_sg2,
yfx,
origin=grid.compute_y_origin(),
domain=grid.domain_x_compute_ybuffer(),
)
# {
# # TODO for if we end up using the Allreduce and compute cmax globally
# (or locally). For now, hardcoded.
# split = int(grid.npz / 6)
# cmax_stencil1(
# cxd, cyd, cmax, origin=grid.compute_origin(),
# domain=(grid.nic, grid.njc, split)
# )
# cmax_stencil2(
# cxd,
# cyd,
# grid.sin_sg5,
# cmax,
# origin=(grid.is_, grid.js, split),
# domain=(grid.nic, grid.njc, grid.npz - split + 1),
# )
# cmax_flat = np.amax(cmax, axis=(0, 1))
# # cmax_flat is a gt4py storage still, but of dimension [npz+1]...
# cmax_max_all_ranks = cmax_flat.data
# # TODO mpi allreduce.... can we not?
# # comm.Allreduce(cmax_flat, cmax_max_all_ranks, op=MPI.MAX)
# }
cmax_max_all_ranks = 2.0
nsplt = math.floor(1.0 + cmax_max_all_ranks)
# NOTE: cmax is not usually a single value, it varies with k, if return to
# that, make nsplt a column as well and compute frac inside cmax_split_vars.
# nsplt3d = utils.make_storage_from_shape(cyd.shape, origin=grid.compute_origin())
# nsplt3d[:] = nsplt
frac = 1.0
if nsplt > 1.0:
frac = 1.0 / nsplt
cmax_multiply_by_frac(
cxd,
xfx,
mfxd,
cyd,
yfx,
mfyd,
frac,
origin=grid.default_origin(),
domain=grid.domain_shape_buffer_1cell(),
)
# complete HALO update on q
for qname in utils.tracer_variables[0:nq]:
q = tracers[qname + "_quantity"]
comm.halo_update(q, n_points=utils.halo)
ra_x_stencil(
grid.area,
xfx,
ra_x,
origin=grid.compute_x_origin(),
domain=grid.domain_y_compute_x(),
)
ra_y_stencil(
grid.area,
yfx,
ra_y,
origin=grid.compute_y_origin(),
domain=grid.domain_x_compute_y(),
)
# TODO: Revisit: the loops over q and nsplt have two inefficient options
# duplicating storages/stencil calls, return to this, maybe you have more
# options now, or maybe the one chosen here is the worse one.
dp1_orig = copy(
dp1, origin=grid.default_origin(), domain=grid.domain_shape_standard()
)
for qname in utils.tracer_variables[0:nq]:
q = tracers[qname + "_quantity"]
# handling the q and it loop switching
copy_stencil(
dp1_orig,
dp1,
origin=grid.default_origin(),
domain=grid.domain_shape_standard(),
)
for it in range(int(nsplt)):
dp_fluxadjustment(
dp1,
mfxd,
mfyd,
grid.rarea,
dp2,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
if nsplt != 1:
if it == 0:
# TODO 1d
qn2 = grid.quantity_wrap(
copy(
q.storage,
origin=grid.default_origin(),
domain=grid.domain_shape_standard(),
),
units="kg/m^2",
)
fvtp2d.compute_no_sg(
qn2.storage,
cxd,
cyd,
spec.namelist.hord_tr,
xfx,
yfx,
ra_x,
ra_y,
fx,
fy,
mfx=mfxd,
mfy=mfyd,
)
if it < nsplt - 1:
q_adjust(
qn2.storage,
dp1,
fx,
fy,
grid.rarea,
dp2,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
else:
q_other_adjust(
qn2.storage,
q.storage,
dp1,
fx,
fy,
grid.rarea,
dp2,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
else:
fvtp2d.compute_no_sg(
q.storage,
cxd,
cyd,
spec.namelist.hord_tr,
xfx,
yfx,
ra_x,
ra_y,
fx,
fy,
mfx=mfxd,
mfy=mfyd,
)
q_adjust(
q.storage,
dp1,
fx,
fy,
grid.rarea,
dp2,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
if it < nsplt - 1:
copy_stencil(
dp2,
dp1,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
comm.halo_update(qn2, n_points=utils.halo)
def compute(state, comm):
# u, v, w, delz, delp, pt, pe, pk, phis, wsd, omga, ua, va, uc, vc, mfxd,
# mfyd, cxd, cyd, pkz, peln, q_con, ak, bk, diss_estd, cappa, mdt, n_split,
# akap, ptop, pfull, n_map, comm):
grid = spec.grid
init_step = state.n_map == 1
end_step = state.n_map == spec.namelist.k_split
akap = state.akap
# peln1 = math.log(ptop)
# ptk = ptop**akap
dt = state.mdt / state.n_split
dt2 = 0.5 * dt
hydrostatic = spec.namelist.hydrostatic
rgrav = 1.0 / constants.GRAV
n_split = state.n_split
# TODO: Put defaults into code.
# m_split = 1. + abs(dt_atmos)/real(k_split*n_split*abs(p_split))
# n_split = nint( real(n0split)/real(k_split*abs(p_split)) * stretch_fac + 0.5 )
ms = max(1, spec.namelist.m_split / 2.0)
shape = state.delz.shape
# NOTE: In Fortran model the halo update starts happens in fv_dynamics, not here.
reqs = {}
for halovar in [
"q_con_quantity", "cappa_quantity", "delp_quantity", "pt_quantity"
]:
reqs[halovar] = comm.start_halo_update(state.__getattribute__(halovar),
n_points=utils.halo)
reqs_vector = comm.start_vector_halo_update(state.u_quantity,
state.v_quantity,
n_points=utils.halo)
reqs["q_con_quantity"].wait()
reqs["cappa_quantity"].wait()
state.__dict__.update(dyncore_temporaries(shape))
if init_step:
state.gz[:-1, :-1, :] = HUGE_R
state.diss_estd[grid.slice_dict(grid.compute_dict())] = 0.0
if not hydrostatic:
state.pk3[:-1, :-1, :] = HUGE_R
state.mfxd[grid.slice_dict(grid.x3d_compute_dict())] = 0.0
state.mfyd[grid.slice_dict(grid.y3d_compute_dict())] = 0.0
state.cxd[grid.slice_dict(grid.x3d_compute_domain_y_dict())] = 0.0
state.cyd[grid.slice_dict(grid.y3d_compute_domain_x_dict())] = 0.0
if not hydrostatic:
# k1k = akap / (1.0 - akap)
# TODO: Is really just a column... when different shapes are supported
# perhaps change this.
state.dp_ref = utils.make_storage_from_shape(state.ak.shape,
grid.default_origin())
state.zs = utils.make_storage_from_shape(state.ak.shape,
grid.default_origin())
dp_ref_compute(
state.ak,
state.bk,
state.phis,
state.dp_ref,
state.zs,
rgrav,
origin=grid.default_origin(),
domain=grid.domain_shape_standard(add=(0, 0, 1)),
)
n_con = get_n_con()
for it in range(n_split):
remap_step = False
if spec.namelist.breed_vortex_inline or (it == n_split - 1):
remap_step = True
if not hydrostatic:
reqs["w_quantity"] = comm.start_halo_update(state.w_quantity,
n_points=utils.halo)
if it == 0:
set_gz(
state.zs,
state.delz,
state.gz,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, grid.npz + 1),
)
reqs["gz_quantity"] = comm.start_halo_update(
state.gz_quantity, n_points=utils.halo)
if it == 0:
reqs["delp_quantity"].wait()
reqs["pt_quantity"].wait()
beta_d = 0
else:
beta_d = spec.namelist.beta
last_step = False
if it == n_split - 1 and end_step:
last_step = True
if it == n_split - 1 and end_step:
if spec.namelist.use_old_omega: # apparently True
set_pem(
state.delp,
state.pem,
state.ptop,
origin=(grid.is_ - 1, grid.js - 1, 0),
domain=(grid.nic + 2, grid.njc + 2, grid.npz),
)
reqs_vector.wait()
if not hydrostatic:
reqs["w_quantity"].wait()
state.delpc, state.ptc = c_sw.compute(
state.delp,
state.pt,
state.u,
state.v,
state.w,
state.uc,
state.vc,
state.ua,
state.va,
state.ut,
state.vt,
state.divgd,
state.omga,
dt2,
)
if spec.namelist.nord > 0:
reqs["divgd_quantity"] = comm.start_halo_update(
state.divgd_quantity, n_points=utils.halo)
if not hydrostatic:
if it == 0:
reqs["gz_quantity"].wait()
copy_stencil(
state.gz,
state.zh,
origin=grid.default_origin(),
domain=grid.domain_shape_buffer_k(),
)
else:
copy_stencil(
state.gz,
state.zh,
origin=grid.default_origin(),
domain=grid.domain_shape_buffer_k(),
)
if not hydrostatic:
state.gz, state.ws3 = updatedzc.compute(state.dp_ref, state.zs,
state.ut, state.vt,
state.gz, state.ws3, dt2)
# TODO: This is really a 2d field.
state.ws3 = utils.make_storage_data(state.ws3[:, :, -1],
shape,
origin=(0, 0, 0))
riem_solver_c.compute(
ms,
dt2,
akap,
state.cappa,
state.ptop,
state.phis,
state.omga,
state.ptc,
state.q_con,
state.delpc,
state.gz,
state.pkc,
state.ws3,
)
pgradc.compute(state.uc, state.vc, state.delpc, state.pkc, state.gz,
dt2)
reqc_vector = comm.start_vector_halo_update(state.uc_quantity,
state.vc_quantity,
n_points=utils.halo)
if spec.namelist.nord > 0:
reqs["divgd_quantity"].wait()
reqc_vector.wait()
state.nord_v, state.damp_vt = d_sw.compute(
state.vt,
state.delp,
state.ptc,
state.pt,
state.u,
state.v,
state.w,
state.uc,
state.vc,
state.ua,
state.va,
state.divgd,
state.mfxd,
state.mfyd,
state.cxd,
state.cyd,
state.crx,
state.cry,
state.xfx,
state.yfx,
state.q_con,
state.zh,
state.heat_source,
state.diss_estd,
dt,
)
for halovar in ["delp_quantity", "pt_quantity", "q_con_quantity"]:
comm.halo_update(state.__getattribute__(halovar),
n_points=utils.halo)
# Not used unless we implement other betas and alternatives to nh_p_grad
# if spec.namelist.d_ext > 0:
# raise 'Unimplemented namelist option d_ext > 0'
# else:
# divg2 = utils.make_storage_from_shape(delz.shape, grid.compute_origin())
if not hydrostatic:
updatedzd.compute(
state.nord_v,
state.damp_vt,
state.dp_ref,
state.zs,
state.zh,
state.crx,
state.cry,
state.xfx,
state.yfx,
state.wsd,
dt,
)
# TODO: This is really a 2d field.
state.wsd = utils.make_storage_data(state.wsd[:, :, -1],
shape,
origin=grid.compute_origin())
riem_solver3.compute(
remap_step,
dt,
akap,
state.cappa,
state.ptop,
state.zs,
state.w,
state.delz,
state.q_con,
state.delp,
state.pt,
state.zh,
state.pe,
state.pkc,
state.pk3,
state.pk,
state.peln,
state.wsd,
)
reqs["zh_quantity"] = comm.start_halo_update(state.zh_quantity,
n_points=utils.halo)
if grid.npx == grid.npy:
reqs["pkc_quantity"] = comm.start_halo_update(
state.pkc_quantity, n_points=2)
else:
reqs["pkc_quantity"] = comm.start_halo_update(
state.pkc_quantity, n_points=utils.halo)
if remap_step:
pe_halo.compute(state.pe, state.delp, state.ptop)
if spec.namelist.use_logp:
raise Exception("unimplemented namelist option use_logp=True")
else:
pk3_halo.compute(state.pk3, state.delp, state.ptop, akap)
if not hydrostatic:
reqs["zh_quantity"].wait()
if grid.npx != grid.npy:
reqs["pkc_quantity"].wait()
if not hydrostatic:
basic.multiply_constant(
state.zh,
state.gz,
constants.GRAV,
origin=(grid.is_ - 2, grid.js - 2, 0),
domain=(grid.nic + 4, grid.njc + 4, grid.npz + 1),
)
if grid.npx == grid.npy:
reqs["pkc_quantity"].wait()
if spec.namelist.beta != 0:
raise Exception(
"Unimplemented namelist option -- we only support beta=0")
if not hydrostatic:
nh_p_grad.compute(
state.u,
state.v,
state.pkc,
state.gz,
state.pk3,
state.delp,
dt,
state.ptop,
akap,
)
if spec.namelist.rf_fast:
# TODO: Pass through ks, or remove, inconsistent representation vs Fortran.
ray_fast.compute(
state.u,
state.v,
state.w,
state.dp_ref,
state.pfull,
dt,
state.ptop,
state.ks,
)
if it != n_split - 1:
reqs_vector = comm.start_vector_halo_update(state.u_quantity,
state.v_quantity,
n_points=utils.halo)
else:
if spec.namelist.grid_type < 4:
comm.synchronize_vector_interfaces(state.u_quantity,
state.v_quantity)
if n_con != 0 and spec.namelist.d_con > 1.0e-5:
nf_ke = min(3, spec.namelist.nord + 1)
comm.halo_update(state.heat_source_quantity, n_points=utils.halo)
cd = constants.CNST_0P20 * grid.da_min
del2cubed.compute(state.heat_source, nf_ke, cd, grid.npz)
if not hydrostatic:
temperature_adjust.compute(
state.pt,
state.pkz,
state.heat_source,
state.delz,
state.delp,
state.cappa,
n_con,
dt,
)
def compute(
u,
v,
va,
ptc,
vort,
ua,
divg_d,
vc,
uc,
delpc,
ke,
wk,
d2_bg,
dt,
nord,
kstart=0,
nk=None,
):
grid = spec.grid
if nk is None:
nk = grid.npz - kstart
# Avoid running center-domain computation on tile edges, since they'll be
# overwritten.
is2 = grid.is_ + 1 if grid.west_edge else grid.is_
ie1 = grid.ie if grid.east_edge else grid.ie + 1
nord = int(nord)
if nord == 0:
damping_zero_order(u, v, va, ptc, vort, ua, vc, uc, delpc, ke, d2_bg,
dt, is2, ie1, kstart, nk)
else:
copy_stencil(
divg_d,
delpc,
origin=(grid.is_, grid.js, kstart),
domain=(grid.nic + 1, grid.njc + 1, nk),
)
for n in range(1, nord + 1):
nt = nord - n
nint = grid.nic + 2 * nt + 1
njnt = grid.njc + 2 * nt + 1
js = grid.js - nt
is_ = grid.is_ - nt
fillc = ((n != nord) and spec.namelist.grid_type < 3
and not grid.nested
and (grid.sw_corner or grid.se_corner or grid.ne_corner
or grid.nw_corner))
if fillc:
corners.fill_corners_2d(divg_d, grid, "B", "x")
vc_from_divg(
divg_d,
grid.divg_u,
vc,
origin=(is_ - 1, js, kstart),
domain=(nint + 1, njnt, nk),
)
if fillc:
corners.fill_corners_2d(divg_d, grid, "B", "y")
uc_from_divg(
divg_d,
grid.divg_v,
uc,
origin=(is_, js - 1, kstart),
domain=(nint, njnt + 1, nk),
)
if fillc:
corners.fill_corners_dgrid(vc, uc, grid, True)
redo_divg_d(uc,
vc,
divg_d,
origin=(is_, js, kstart),
domain=(nint, njnt, nk))
corner_domain = (1, 1, nk)
if grid.sw_corner:
corner_south_remove_extra_term(uc,
divg_d,
origin=(grid.is_, grid.js,
kstart),
domain=corner_domain)
if grid.se_corner:
corner_south_remove_extra_term(
uc,
divg_d,
origin=(grid.ie + 1, grid.js, kstart),
domain=corner_domain,
)
if grid.ne_corner:
corner_north_remove_extra_term(
uc,
divg_d,
origin=(grid.ie + 1, grid.je + 1, kstart),
domain=corner_domain,
)
if grid.nw_corner:
corner_north_remove_extra_term(
uc,
divg_d,
origin=(grid.is_, grid.je + 1, kstart),
domain=corner_domain,
)
if not grid.stretched_grid:
basic.adjustmentfactor_stencil(
grid.rarea_c,
divg_d,
origin=(is_, js, kstart),
domain=(nint, njnt, nk),
)
vorticity_calc(wk, vort, delpc, dt, nord, kstart, nk)
if grid.stretched_grid:
dd8 = grid.da_min * spec.namelist.d4_bg**(nord + 1)
else:
dd8 = (grid.da_min_c * spec.namelist.d4_bg)**(nord + 1)
damping_nord_highorder_stencil(
vort,
ke,
delpc,
divg_d,
grid.da_min_c,
d2_bg,
spec.namelist.dddmp,
dd8,
origin=(grid.is_, grid.js, kstart),
domain=(grid.nic + 1, grid.njc + 1, nk),
)
return vort, ke, delpc
def compute(
tracers,
pt,
delp,
delz,
peln,
u,
v,
w,
ua,
cappa,
q_con,
pkz,
pk,
pe,
hs,
te,
ps,
wsd,
omga,
ak,
bk,
gz,
cvm,
ptop,
akap,
r_vir,
nq,
):
grid = spec.grid
hydrostatic = spec.namelist.hydrostatic
t_min = 184.0
# do_omega = hydrostatic and last_step # TODO pull into inputs
domain_jextra = (grid.nic, grid.njc + 1, grid.npz + 1)
pe1 = copy(pe, origin=grid.compute_origin(), domain=domain_jextra)
pe2 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
dp2 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
pn2 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
pe0 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
pe3 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
# pk2 = utils.make_storage_from_shape(pe.shape, grid.compute_origin())
init_pe2(pe, pe2, ptop, origin=grid.compute_origin(), domain=domain_jextra)
if spec.namelist.kord_tm < 0:
if hydrostatic:
raise Exception("Hydrostatic is not implemented")
else:
moist_cv.compute_pt(
tracers["qvapor"],
tracers["qliquid"],
tracers["qice"],
tracers["qrain"],
tracers["qsnow"],
tracers["qgraupel"],
q_con,
gz,
cvm,
pt,
cappa,
delp,
delz,
r_vir,
)
if not hydrostatic:
delz_adjust(delp,
delz,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute())
pressure_updates(
pe1,
pe2,
pe,
ak,
bk,
dp2,
ps,
pn2,
peln,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute_buffer_k(),
)
# TODO: Fix silly hack due to pe2 being 2d, so pe[:, je+1, 1:npz] should be
# the same as it was for pe[:, je, 1:npz] (unchanged)
copy_j_adjacent(pe2,
origin=(grid.is_, grid.je + 1, 1),
domain=(grid.nic, 1, grid.npz - 1))
copy_stencil(dp2,
delp,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute())
pn2_and_pk(
pe2,
pn2,
pk,
akap,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
if spec.namelist.kord_tm < 0:
map_single.compute(
pt,
peln,
pn2,
gz,
1,
grid.is_,
grid.ie,
abs(spec.namelist.kord_tm),
qmin=t_min,
)
else:
raise Exception("map ppm, untested mode where kord_tm >= 0")
map_single.compute(
pt,
pe1,
pe2,
gz,
1,
grid.is_,
grid.ie,
abs(spec.namelist.kord_tm),
qmin=t_min,
)
# TODO if nq > 5:
mapn_tracer.compute(pe1, pe2, dp2, tracers, nq, 0.0, grid.is_, grid.ie,
abs(spec.namelist.kord_tr))
# TODO else if nq > 0:
# TODO map1_q2, fillz
if not hydrostatic:
map_single.compute(w, pe1, pe2, wsd, -2, grid.is_, grid.ie,
spec.namelist.kord_wz)
map_single.compute(delz, pe1, pe2, gz, 1, grid.is_, grid.ie,
spec.namelist.kord_wz)
undo_delz_adjust(delp,
delz,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute())
# if do_omega: # NOTE untested
# pe3 = copy(omga, origin=(grid.is_, grid.js, 1))
pe0 = copy(peln,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, grid.npz + 1))
copy_stencil(
pn2,
peln,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, grid.npz + 1),
)
if hydrostatic:
# pkz
pass
else:
moist_cv.compute_pkz(
tracers["qvapor"],
tracers["qliquid"],
tracers["qice"],
tracers["qrain"],
tracers["qsnow"],
tracers["qgraupel"],
q_con,
gz,
cvm,
pkz,
pt,
cappa,
delp,
delz,
r_vir,
)
# if do_omega:
# dp2 update, if larger than pe0 and smaller than one level up, update omega
# and exit
pressures_mapu(pe,
pe1,
ak,
bk,
pe0,
pe3,
origin=grid.compute_origin(),
domain=domain_jextra)
map_single.compute(u,
pe0,
pe3,
gz,
-1,
grid.is_,
grid.ie,
spec.namelist.kord_mt,
j_interface=True)
domain_iextra = (grid.nic + 1, grid.njc, grid.npz + 1)
pressures_mapv(pe,
ak,
bk,
pe0,
pe3,
origin=grid.compute_origin(),
domain=domain_iextra)
map_single.compute(v, pe0, pe3, gz, -1, grid.is_, grid.ie + 1,
spec.namelist.kord_mt)
update_ua(pe2, ua, origin=grid.compute_origin(), domain=domain_jextra)