def compute(ms, dt2, akap, cappa, ptop, hs, w3, ptc, q_con, delpc, gz, pef,
ws):
grid = spec.grid
is1 = grid.is_ - 1
ie1 = grid.ie + 1
js1 = grid.js - 1
je1 = grid.je + 1
km = spec.grid.npz - 1
islice = slice(is1, ie1 + 1)
kslice = slice(0, km + 1)
kslice_shift = slice(1, km + 2)
shape = w3.shape
domain = (spec.grid.nic + 2, grid.njc + 2, km + 2)
riemorigin = (is1, js1, 0)
dm = copy(delpc)
cp3 = copy(cappa)
w = copy(w3)
pem = utils.make_storage_from_shape(shape, riemorigin)
peg = utils.make_storage_from_shape(shape, riemorigin)
pe = utils.make_storage_from_shape(shape, riemorigin)
gm = utils.make_storage_from_shape(shape, riemorigin)
dz = utils.make_storage_from_shape(shape, riemorigin)
pm = utils.make_storage_from_shape(shape, riemorigin)
precompute(
cp3,
gz,
dm,
q_con,
pem,
peg,
dz,
gm,
pef,
pm,
ptop,
origin=riemorigin,
domain=domain,
)
sim1_solver.solve(is1, ie1, js1, je1, dt2, gm, cp3, pe, dm, pm, pem, w, dz,
ptc, ws)
finalize(pe, pem, hs, dz, pef, gz, origin=riemorigin, domain=domain)
def setup_data(
q1: FloatField,
pe1: FloatField,
i1: int,
i2: int,
j_2d: Optional[int] = None,
j_interface: bool = False,
):
grid = spec.grid
i_extent = i2 - i1 + 1
origin, domain, jslice = region_mode(j_2d, i1, i_extent, grid)
if j_interface:
jslice = slice(jslice.start, jslice.stop + 1)
domain = (domain[0], jslice.stop - jslice.start, domain[2])
dp1 = utils.make_storage_from_shape(q1.shape, origin=origin)
q4_1 = copy(q1, origin=(0, 0, 0), domain=grid.domain_shape_standard())
q4_2 = utils.make_storage_from_shape(q4_1.shape, origin=(grid.is_, 0, 0))
q4_3 = utils.make_storage_from_shape(q4_1.shape, origin=(grid.is_, 0, 0))
q4_4 = utils.make_storage_from_shape(q4_1.shape, origin=(grid.is_, 0, 0))
set_dp(dp1, pe1, origin=origin, domain=domain)
return dp1, q4_1, q4_2, q4_3, q4_4, origin, domain, jslice, i_extent
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 column_calls(
zh, crx_adv, cry_adv, xfx_adv, yfx_adv, ra_x, ra_y, ndif, damp, kstart, nk
):
grid = spec.grid
default_origin = (grid.isd, grid.jsd, kstart)
compute_origin = (grid.is_, grid.js, kstart)
compute_domain = (grid.nic, grid.njc, nk)
if damp > 1e-5:
wk = utils.make_storage_from_shape(zh.shape, default_origin)
fx2 = utils.make_storage_from_shape(zh.shape, default_origin)
fy2 = utils.make_storage_from_shape(zh.shape, default_origin)
fx = utils.make_storage_from_shape(zh.shape, default_origin)
fy = utils.make_storage_from_shape(zh.shape, default_origin)
z2 = copy(zh, origin=default_origin, domain=(grid.nid, grid.njd, nk))
fvtp2d.compute_no_sg(
z2,
crx_adv,
cry_adv,
spec.namelist.hord_tm,
xfx_adv,
yfx_adv,
ra_x,
ra_y,
fx,
fy,
kstart=kstart,
nk=nk,
)
delnflux.compute_no_sg(z2, fx2, fy2, ndif, damp, wk, kstart=kstart, nk=nk)
zh_damp_stencil(
grid.area,
z2,
fx,
fy,
ra_x,
ra_y,
fx2,
fy2,
grid.rarea,
zh,
origin=compute_origin,
domain=compute_domain,
)
else:
raise Exception("untested")
fvtp2d.compute_no_sg(
zh,
crx_adv,
cry_adv,
spec.namelist.hord_tm,
xfx_adv,
yfx_adv,
ra_x,
ra_y,
fx,
fy,
kstart=kstart,
nk=nk,
)
zh_stencil(
grid.area,
zh,
fx,
fy,
ra_x,
ra_y,
origin=compute_origin,
domain=compute_domain,
)
return [zh]
def compute_no_sg(q, fx2, fy2, nord, damp_c, d2, kstart=0, nk=None, mass=None):
grid = spec.grid
nord = int(nord)
i1 = grid.is_ - 1 - nord
i2 = grid.ie + 1 + nord
j1 = grid.js - 1 - nord
j2 = grid.je + 1 + nord
if nk is None:
nk = grid.npz - kstart
kslice = slice(kstart, kstart + nk)
origin_d2 = (i1, j1, kstart)
domain_d2 = (i2 - i1 + 1, j2 - j1 + 1, nk)
if mass is None:
d2_damp(q, d2, damp_c, origin=origin_d2, domain=domain_d2)
else:
d2 = copy(q, origin=origin_d2, domain=domain_d2)
if nord > 0:
corners.copy_corners(d2, "x", grid, kslice)
f1_ny = grid.je - grid.js + 1 + 2 * nord
f1_nx = grid.ie - grid.is_ + 2 + 2 * nord
fx_origin = (grid.is_ - nord, grid.js - nord, kstart)
fx2_order(d2,
grid.del6_v,
fx2,
order=1,
origin=fx_origin,
domain=(f1_nx, f1_ny, nk))
if nord > 0:
corners.copy_corners(d2, "y", grid, kslice)
fy2_order(
d2,
grid.del6_u,
fy2,
order=1,
origin=fx_origin,
domain=(f1_nx - 1, f1_ny + 1, nk),
)
if nord > 0:
for n in range(nord):
nt = nord - 1 - n
nt_origin = (grid.is_ - nt - 1, grid.js - nt - 1, kstart)
nt_ny = grid.je - grid.js + 3 + 2 * nt
nt_nx = grid.ie - grid.is_ + 3 + 2 * nt
d2_highorder(fx2,
fy2,
grid.rarea,
d2,
origin=nt_origin,
domain=(nt_nx, nt_ny, nk))
corners.copy_corners(d2, "x", grid, kslice)
nt_origin = (grid.is_ - nt, grid.js - nt, kstart)
fx2_order(
d2,
grid.del6_v,
fx2,
order=2 + n,
origin=nt_origin,
domain=(nt_nx - 1, nt_ny - 2, nk),
)
corners.copy_corners(d2, "y", grid, kslice)
fy2_order(
d2,
grid.del6_u,
fy2,
order=2 + n,
origin=nt_origin,
domain=(nt_nx - 2, nt_ny - 1, nk),
)
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(
tracers: Dict[str, Any],
pt: FloatField,
delp: FloatField,
delz: FloatField,
peln: FloatField,
u: FloatField,
v: FloatField,
w: FloatField,
ua: FloatField,
cappa: FloatField,
q_con: FloatField,
pkz: FloatField,
pk: FloatField,
pe: FloatField,
hs: FloatField,
te: FloatField,
ps: FloatField,
wsd: FloatField,
omga: FloatField,
ak: FloatField,
bk: FloatField,
gz: FloatField,
cvm: FloatField,
ptop: float,
akap: float,
r_vir: float,
nq: int,
):
grid = spec.grid
hydrostatic = spec.namelist.hydrostatic
t_min = 184.0
if spec.namelist.kord_tm >= 0:
raise Exception("map ppm, untested mode where kord_tm >= 0")
if hydrostatic:
raise Exception("Hydrostatic is not implemented")
# 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)
moist_cv_pt_pressure(
tracers["qvapor"],
tracers["qliquid"],
tracers["qrain"],
tracers["qsnow"],
tracers["qice"],
tracers["qgraupel"],
q_con,
gz,
cvm,
pt,
cappa,
delp,
delz,
pe,
pe2,
ak,
bk,
dp2,
ps,
pn2,
peln,
r_vir,
hydrostatic,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute_buffer_k(),
)
pn2_and_pk(
pe2,
pn2,
pk,
akap,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(add=(0, 1, 0)),
)
map_single.compute(
pt,
peln,
pn2,
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
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_and_copy_peln(
delp,
delz,
peln,
pe0,
pn2,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, grid.npz + 1),
)
# if do_omega: # NOTE untested
# pe3 = copy(omga, origin=(grid.is_, grid.js, 1))
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)
def compute(
last_call,
dt,
akap,
cappa,
ptop,
zs,
w,
delz,
q_con,
delp,
pt,
zh,
pe,
ppe,
pk3,
pk,
peln,
wsd,
):
grid = spec.grid
rgrav = 1.0 / constants.GRAV
km = grid.npz - 1
peln1 = math.log(ptop)
ptk = math.exp(akap * peln1)
islice = slice(grid.is_, grid.ie + 1)
kslice = slice(0, km + 1)
kslice_shift = slice(1, km + 2)
shape = w.shape
domain = (grid.nic, grid.njc, km + 2)
riemorigin = (grid.is_, grid.js, 0)
dm = copy(delp)
cp3 = copy(cappa)
pe_init = copy(pe)
pm = utils.make_storage_from_shape(shape, riemorigin)
pem = utils.make_storage_from_shape(shape, riemorigin)
peln_run = utils.make_storage_from_shape(shape, riemorigin)
peg = utils.make_storage_from_shape(shape, riemorigin)
pelng = utils.make_storage_from_shape(shape, riemorigin)
gm = utils.make_storage_from_shape(shape, riemorigin)
precompute(
cp3,
dm,
zh,
q_con,
pem,
peln_run,
pk3,
peg,
pelng,
gm,
delz,
pm,
ptop,
peln1,
ptk,
rgrav,
akap,
origin=riemorigin,
domain=domain,
)
sim1_solver.solve(
grid.is_,
grid.ie,
grid.js,
grid.je,
dt,
gm,
cp3,
pe,
dm,
pm,
pem,
w,
delz,
pt,
wsd,
)
finalize(
zs,
delz,
zh,
peln_run,
peln,
pk3,
pk,
pem,
pe,
ppe,
pe_init,
last_call,
origin=riemorigin,
domain=domain,
)
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)