def _initialize_interpolation_constants(self,
stencil_factory: StencilFactory,
grid_indexing: GridIndexing):
# because stencils only work on 3D at the moment, need to compute in 3D
# and then make these 1D
gk_3d = utils.make_storage_from_shape(
(1, 1, grid_indexing.domain[2] + 1), (0, 0, 0))
gamma_3d = utils.make_storage_from_shape(
(1, 1, grid_indexing.domain[2] + 1), (0, 0, 0))
beta_3d = utils.make_storage_from_shape(
(1, 1, grid_indexing.domain[2] + 1), (0, 0, 0))
_cubic_spline_interpolation_constants = stencil_factory.from_origin_domain(
cubic_spline_interpolation_constants,
origin=(0, 0, 0),
domain=(1, 1, grid_indexing.domain[2] + 1),
)
_cubic_spline_interpolation_constants(self._dp0, gk_3d, beta_3d,
gamma_3d)
self._gk = utils.make_storage_data(gk_3d[0, 0, :], gk_3d.shape[2:],
(0, ))
self._beta = utils.make_storage_data(beta_3d[0, 0, :],
beta_3d.shape[2:], (0, ))
self._gamma = utils.make_storage_data(gamma_3d[0, 0, :],
gamma_3d.shape[2:], (0, ))
self._copy_corners = PreAllocatedCopiedCornersFactory(
stencil_factory=stencil_factory,
dims=[
fv3gfs.util.X_DIM, fv3gfs.util.Y_DIM,
fv3gfs.util.Z_INTERFACE_DIM
],
y_temporary=None,
)
def compute_from_storage(self, inputs):
inputs["q_x_flux"] = utils.make_storage_from_shape(
self.maxshape, self.grid.full_origin()
)
inputs["q_y_flux"] = utils.make_storage_from_shape(
self.maxshape, self.grid.full_origin()
)
for optional_arg in ["mass"]:
if optional_arg not in inputs:
inputs[optional_arg] = None
self.compute_func = FiniteVolumeTransport(
stencil_factory=self.grid.stencil_factory,
grid_data=self.grid.grid_data,
damping_coefficients=self.grid.damping_coefficients,
grid_type=self.grid.grid_type,
hord=int(inputs["hord"]),
nord=inputs.pop("nord"),
damp_c=inputs.pop("damp_c"),
)
del inputs["hord"]
q_storage = inputs["q"]
factory = PreAllocatedCopiedCornersFactory(
self.grid.stencil_factory, dims=[X_DIM, Y_DIM, Z_DIM], y_temporary=None
)
inputs["q"] = factory(q_storage)
self.compute_func(**inputs)
inputs["q"] = q_storage
return inputs
def __init__(
self,
stencil_factory: StencilFactory,
im: int,
jm: int,
km: int,
nq: int,
):
self._nq = nq
self._fix_tracer_stencil = stencil_factory.from_origin_domain(
fix_tracer,
origin=stencil_factory.grid_indexing.origin_compute(),
domain=(im, jm, km),
)
shape = stencil_factory.grid_indexing.domain_full(add=(1, 1, 1))
shape_ij = shape[0:2]
self._dm = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
self._dm_pos = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
# Setting initial value of upper_fix to zero is only needed for validation.
# The values in the compute domain are set to zero in the stencil.
self._zfix = utils.make_storage_from_shape(shape_ij,
dtype=int,
origin=(0, 0))
self._sum0 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
self._sum1 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
def compute(uc: sd, vc: sd, u: sd, v: sd, ua: sd, va: sd, dt2: float):
grid = spec.grid
origin = (grid.is_ - 1, grid.js - 1, 0)
# Create storage objects to hold the new vorticity and kinetic energy values
ke_c = utils.make_storage_from_shape(uc.shape, origin=origin)
vort_c = utils.make_storage_from_shape(vc.shape, origin=origin)
# Set vorticity and kinetic energy values
update_vorticity_and_kinetic_energy(
ke_c,
vort_c,
ua,
va,
uc,
vc,
u,
v,
grid.sin_sg1,
grid.cos_sg1,
grid.sin_sg2,
grid.cos_sg2,
grid.sin_sg3,
grid.cos_sg3,
grid.sin_sg4,
grid.cos_sg4,
dt2,
origin=(grid.is_ - 1, grid.js - 1, 0),
domain=(grid.nic + 2, grid.njc + 2, grid.npz),
)
return ke_c, vort_c
def compute(u, v, pp, gz, pk3, delp, dt, ptop, akap):
"""
u=u v=v pp=pkc gz=gz pk3=pk3 delp=delp dt=dt
"""
grid = spec.grid
orig = (grid.is_, grid.js, 0)
# peln1 = log(ptop)
ptk = ptop**akap
top_value = ptk # = peln1 if spec.namelist.use_logp else ptk
wk1 = utils.make_storage_from_shape(pp.shape, origin=orig)
wk = utils.make_storage_from_shape(pk3.shape, origin=orig)
set_k0(pp,
pk3,
top_value,
origin=orig,
domain=(grid.nic + 1, grid.njc + 1, 1))
a2b_ord4.compute(pp, wk1, kstart=1, nk=grid.npz, replace=True)
a2b_ord4.compute(pk3, wk1, kstart=1, nk=grid.npz, replace=True)
a2b_ord4.compute(gz, wk1, kstart=0, nk=grid.npz + 1, replace=True)
a2b_ord4.compute(delp, wk1)
CalcWk(pk3, wk, origin=orig, domain=(grid.nic + 1, grid.njc + 1, grid.npz))
du = utils.make_storage_from_shape(u.shape, origin=orig)
CalcU(
u,
du,
wk,
wk1,
gz,
pk3,
pp,
grid.rdx,
dt,
origin=orig,
domain=(grid.nic, grid.njc + 1, grid.npz),
)
dv = utils.make_storage_from_shape(v.shape, origin=orig)
CalcV(
v,
dv,
wk,
wk1,
gz,
pk3,
pp,
grid.rdy,
dt,
origin=orig,
domain=(grid.nic + 1, grid.njc, grid.npz),
)
return u, v, pp, gz, pk3, delp
def _allocate_temporary_storages(self, grid_indexing):
largest_possible_shape = grid_indexing.domain_full(add=(1, 1, 1))
self._crx_interface = utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)),
)
self._cry_interface = utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(-grid_indexing.n_halo, 0, 0)),
)
self._x_area_flux_interface = utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)),
)
self._y_area_flux_interface = utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(-grid_indexing.n_halo, 0, 0)),
)
self._wk = utils.make_storage_from_shape(largest_possible_shape,
grid_indexing.origin_full())
self._height_x_diffusive_flux = utils.make_storage_from_shape(
largest_possible_shape, grid_indexing.origin_full())
self._height_y_diffusive_flux = utils.make_storage_from_shape(
largest_possible_shape, grid_indexing.origin_full())
self._fx = utils.make_storage_from_shape(largest_possible_shape,
grid_indexing.origin_full())
self._fy = utils.make_storage_from_shape(largest_possible_shape,
grid_indexing.origin_full())
def compute(qvapor, qliquid, qrain, qsnow, qice, qgraupel, qcld, pt, delp,
delz, peln):
grid = spec.grid
i_ext = grid.domain_shape_compute()[0]
j_ext = grid.domain_shape_compute()[1]
k_ext = grid.domain_shape_compute()[2]
if spec.namelist.check_negative:
raise Exception("Unimplemented namelist value check_negative=True")
if spec.namelist.hydrostatic:
d0_vap = constants.CP_VAP - constants.C_LIQ
raise Exception("Unimplemented namelist hydrostatic=True")
else:
d0_vap = constants.CV_VAP - constants.C_LIQ
lv00 = constants.HLV - d0_vap * constants.TICE
fix_neg_water(
pt,
delp,
delz,
qvapor,
qliquid,
qrain,
qsnow,
qice,
qgraupel,
lv00,
d0_vap,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
fillq(qgraupel, delp, grid)
fillq(qrain, delp, grid)
# fix_water_vapor(delp, qvapor, origin=grid.compute_origin(),
# domain=grid.domain_shape_compute())
# fix_water_vapor_nonstencil(grid, qvapor, delp)
# fix_water_vapor_bottom(grid, qvapor, delp)
upper_fix = utils.make_storage_from_shape(qvapor.shape, origin=(0, 0, 0))
lower_fix = utils.make_storage_from_shape(qvapor.shape, origin=(0, 0, 0))
bot_dp = delp[:, :, grid.npz - 1]
full_bot_arr = utils.repeat(bot_dp[:, :, np.newaxis], k_ext + 1, axis=2)
dp_bot = utils.make_storage_data(full_bot_arr)
fix_water_vapor_down(
qvapor,
delp,
upper_fix,
lower_fix,
dp_bot,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute(),
)
qvapor[:, :, grid.npz] = upper_fix[:, :, 0]
fix_neg_cloud(delp,
qcld,
origin=grid.compute_origin(),
domain=grid.domain_shape_compute())
def __init__(
self,
stencil_factory: StencilFactory,
damping_coefficients: DampingCoefficients,
rarea,
nord: FloatFieldK,
damp_c: FloatFieldK,
):
"""
nord sets the order of damping to apply:
nord = 0: del-2
nord = 1: del-4
nord = 2: del-6
nord and damp_c define the damping coefficient used in DelnFluxNoSG
"""
self._no_compute = False
if (damp_c <= 1e-4).all():
self._no_compute = True
elif (damp_c[:-1] <= 1e-4).any():
raise NotImplementedError(
"damp_c currently must be always greater than 10^-4 for delnflux"
)
grid_indexing = stencil_factory.grid_indexing
nk = grid_indexing.domain[2]
self._origin = grid_indexing.origin_full()
shape = grid_indexing.max_shape
k_shape = (1, 1, nk)
self._damp_3d = utils.make_storage_from_shape(k_shape)
# fields must be 3d to assign to them
self._fx2 = utils.make_storage_from_shape(shape)
self._fy2 = utils.make_storage_from_shape(shape)
self._d2 = utils.make_storage_from_shape(grid_indexing.domain_full())
damping_factor_calculation = stencil_factory.from_origin_domain(
calc_damp, origin=(0, 0, 0), domain=k_shape
)
self._add_diffusive_stencil = stencil_factory.from_dims_halo(
func=add_diffusive_component,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
)
self._diffusive_damp_stencil = stencil_factory.from_dims_halo(
func=diffusive_damp, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM]
)
damping_factor_calculation(
self._damp_3d, nord, damp_c, damping_coefficients.da_min
)
self._damp = utils.make_storage_data(self._damp_3d[0, 0, :], (nk,), (0,))
self.delnflux_nosg = DelnFluxNoSG(
stencil_factory, damping_coefficients, rarea, nord, nk=nk
)
def __init__(self, stencil_factory: StencilFactory, grid_data: GridData, grid_type):
grid_indexing = stencil_factory.grid_indexing
self.orig = grid_indexing.origin_compute()
domain_full_k = grid_indexing.domain_compute(add=(1, 1, 0))
u_domain = grid_indexing.domain_compute(add=(0, 1, 0))
v_domain = grid_indexing.domain_compute(add=(1, 0, 0))
self.nk = grid_indexing.domain[2]
self._rdx = grid_data.rdx
self._rdy = grid_data.rdy
self._tmp_wk = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=self.orig
) # pk3.shape
self._tmp_wk1 = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=self.orig
) # pp.shape
self._set_k0_and_calc_wk_stencil = stencil_factory.from_origin_domain(
set_k0_and_calc_wk,
origin=self.orig,
domain=domain_full_k,
)
self._calc_u_stencil = stencil_factory.from_origin_domain(
calc_u,
origin=self.orig,
domain=u_domain,
)
self._calc_v_stencil = stencil_factory.from_origin_domain(
calc_v,
origin=self.orig,
domain=v_domain,
)
self.a2b_k1 = AGrid2BGridFourthOrder(
stencil_factory.restrict_vertical(k_start=1),
grid_data,
grid_type,
z_dim=Z_INTERFACE_DIM,
replace=True,
)
self.a2b_kbuffer = AGrid2BGridFourthOrder(
stencil_factory,
grid_data,
grid_type,
z_dim=Z_INTERFACE_DIM,
replace=True,
)
self.a2b_kstandard = AGrid2BGridFourthOrder(
stencil_factory,
grid_data,
grid_type,
replace=False,
)
def compute(self, inputs):
inputs["fx"] = utils.make_storage_from_shape(
self.maxshape, self.grid.default_origin())
inputs["fy"] = utils.make_storage_from_shape(
self.maxshape, self.grid.default_origin())
for optional_arg in ["mass", "mfx", "mfy"]:
if optional_arg not in inputs:
inputs[optional_arg] = None
return self.column_split_compute(inputs, {
"nord": "nord_column",
"damp_c": "damp_c"
})
def compute(u, v, w, dp, pfull, dt, ptop, ks):
grid = spec.grid
namelist = spec.namelist
# The next 3 variables and dm_stencil could be pushed into ray_fast_wind and
# still work, but then recomputing it all twice.
rf_cutoff_nudge = namelist.rf_cutoff + min(100.0, 10.0 * ptop)
# TODO 1D variable
dm = utils.make_storage_from_shape(u.shape, grid.default_origin())
# TODO 1D variable
rf = utils.make_storage_from_shape(u.shape, grid.default_origin())
dm_stencil(
dp,
dm,
pfull,
rf,
dt,
ptop,
rf_cutoff_nudge,
ks,
origin=grid.compute_origin(),
domain=(grid.nic + 1, grid.njc + 1, grid.npz),
)
ray_fast_wind(
u,
rf,
dp,
dm,
pfull,
rf_cutoff_nudge,
ks,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc + 1, grid.npz),
)
ray_fast_wind(
v,
rf,
dp,
dm,
pfull,
rf_cutoff_nudge,
ks,
origin=grid.compute_origin(),
domain=(grid.nic + 1, grid.njc, grid.npz),
)
if not namelist.hydrostatic:
ray_fast_w(
w,
rf,
pfull,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, grid.npz),
)
def compute_delnflux_no_sg(q,
fx,
fy,
nord,
damp_c,
kstart=0,
nk=None,
d2=None,
mass=None):
grid = spec.grid
if nk is None:
nk = grid.npz - kstart
default_origin = (grid.isd, grid.jsd, kstart)
if d2 is None:
d2 = utils.make_storage_from_shape(q.shape, default_origin)
if damp_c <= 1e-4:
return fx, fy
damp = (damp_c * grid.da_min)**(nord + 1)
fx2 = utils.make_storage_from_shape(q.shape, default_origin)
fy2 = utils.make_storage_from_shape(q.shape, default_origin)
compute_no_sg(q, fx2, fy2, nord, damp, d2, kstart, nk, mass)
diffuse_origin = (grid.is_, grid.js, kstart)
diffuse_domain_x = (grid.nic + 1, grid.njc, nk)
diffuse_domain_y = (grid.nic, grid.njc + 1, nk)
if mass is None:
add_diffusive_component(fx,
fx2,
origin=diffuse_origin,
domain=diffuse_domain_x)
add_diffusive_component(fy,
fy2,
origin=diffuse_origin,
domain=diffuse_domain_y)
else:
# TODO: To join these stencils you need to overcompute, making the edges
# 'wrong', but not actually used, separating now for comparison sanity.
# diffusive_damp(fx, fx2, fy, fy2, mass, damp, origin=diffuse_origin,
# domain=(grid.nic + 1, grid.njc + 1, nk))
diffusive_damp_x(fx,
fx2,
mass,
damp,
origin=diffuse_origin,
domain=diffuse_domain_x)
diffusive_damp_y(fy,
fy2,
mass,
damp,
origin=diffuse_origin,
domain=diffuse_domain_y)
return fx, fy
def damp_vertical_wind(w, heat_s, diss_e, dt, column_namelist):
dw = utils.make_storage_from_shape(w.shape, grid().compute_origin())
wk = utils.make_storage_from_shape(w.shape, grid().default_origin())
fx2 = utils.make_storage_from_shape(w.shape, grid().default_origin())
fy2 = utils.make_storage_from_shape(w.shape, grid().default_origin())
if column_namelist["damp_w"] > 1e-5:
dd8 = column_namelist["ke_bg"] * abs(dt)
damp4 = (column_namelist["damp_w"] *
grid().da_min_c)**(column_namelist["nord_w"] + 1)
delnflux.compute_no_sg(w, fx2, fy2, column_namelist["nord_w"], damp4,
wk)
heatdiss.compute(fx2, fy2, w, dd8, dw, heat_s, diss_e)
return dw, wk
def compute(self, inputs):
self.make_storage_data_input_vars(inputs)
orig = (self.grid.is_ - 1, self.grid.js - 1, 0)
inputs["delpc"] = utils.make_storage_from_shape(inputs["delp"].shape,
origin=orig)
inputs["ptc"] = utils.make_storage_from_shape(inputs["pt"].shape,
origin=orig)
c_sw.transportdelp(
**inputs,
rarea=self.grid.rarea,
origin=orig,
domain=self.grid.domain_shape_compute(add=(2, 2, 0)),
)
return self.slice_output(inputs, )
def compute_qyy(qy, qout, kstart, nk):
qyy = utils.make_storage_from_shape(qy.shape,
origin=grid().default_origin())
# avoid running center-domain computation on tile edges, since they'll be
# overwritten.
js = grid().js + 1 if grid().south_edge else grid().js
je = grid().je if grid().north_edge else grid().je + 1
is_ = grid().is_ + 2 if grid().west_edge else grid().is_
ie = grid().ie - 1 if grid().east_edge else grid().ie + 1
dj = je - js + 1
lagrange_interpolation_x(qy,
qyy,
origin=(is_, js, kstart),
domain=(ie - is_ + 1, dj, nk))
if grid().west_edge:
cubic_interpolation_west(qy,
qout,
qyy,
origin=(grid().is_ + 1, js, kstart),
domain=(1, dj, nk))
if grid().east_edge:
cubic_interpolation_east(qy,
qout,
qyy,
origin=(grid().ie, js, kstart),
domain=(1, dj, nk))
return qyy
def compute_qxx(qx, qout, kstart, nk):
qxx = utils.make_storage_from_shape(qx.shape,
origin=grid().default_origin())
# avoid running center-domain computation on tile edges, since they'll be
# overwritten.
js = grid().js + 2 if grid().south_edge else grid().js
je = grid().je - 1 if grid().north_edge else grid().je + 1
is_ = grid().is_ + 1 if grid().west_edge else grid().is_
ie = grid().ie if grid().east_edge else grid().ie + 1
di = ie - is_ + 1
lagrange_interpolation_y(qx,
qxx,
origin=(is_, js, kstart),
domain=(di, je - js + 1, nk))
if grid().south_edge:
cubic_interpolation_south(qx,
qout,
qxx,
origin=(is_, grid().js + 1, kstart),
domain=(di, 1, nk))
if grid().north_edge:
cubic_interpolation_north(qx,
qout,
qxx,
origin=(is_, grid().je, kstart),
domain=(di, 1, nk))
return qxx
def compute(self, inputs):
self.setup(inputs)
inputs["j_2d"] = self.grid.global_to_local_y(
inputs["j_2d"] + TranslateGrid.fpy_model_index_offset)
inputs["i1"] = self.is_
inputs["i2"] = self.ie
inputs["j1"] = inputs["j_2d"]
inputs["j2"] = inputs["j_2d"]
del inputs["j_2d"]
inputs["kord"] = abs(spec.namelist.kord_tm)
inputs["qmin"] = 184.0
# these are sometimes 3D and sometimes singleton in J
if inputs["pe1"].shape[1] == 1:
inputs["pe1"] = self.make_storage_data(
pad_field_in_j(inputs["pe1"], self.nj))
if inputs["pe2"].shape[1] == 1:
inputs["pe2"] = self.make_storage_data(
pad_field_in_j(inputs["pe2"], self.nj))
qs_field = utils.make_storage_from_shape(self.maxshape[0:2],
origin=(0, 0))
qs_field[:, :] = inputs["qs"][:, :, 0]
inputs["qs"] = qs_field
if inputs["qs"].shape[1] == 1:
inputs["qs"] = utils.tile(inputs["qs"][:, 0],
[self.nj, 1]).transpose(1, 0)
var_inout = self.compute_func(**inputs)
return self.slice_output(inputs, {"pt": var_inout})
def compute(qdel, nmax, cd, km):
grid = spec.grid
origin = (grid.isd, grid.jsd, 0)
# Construct some necessary temporary storage objects
fx = utils.make_storage_from_shape(qdel.shape, origin=origin)
fy = utils.make_storage_from_shape(qdel.shape, origin=origin)
# set up the temporal loop
ntimes = min(3, nmax)
for n in range(1, ntimes + 1):
nt = ntimes - n
origin = (grid.is_ - nt, grid.js - nt, 0)
# Fill in appropriate corner values
qdel = corner_fill(grid, qdel)
if nt > 0:
corners.copy_corners(qdel, "x", grid)
nx = grid.njc + 2 * nt + 1 # (grid.ie+nt+1) - (grid.is_-nt) + 1
ny = grid.njc + 2 * nt # (grid.je+nt) - (grid.js-nt) + 1
compute_zonal_flux(fx,
qdel,
grid.del6_v,
origin=origin,
domain=(nx, ny, km))
if nt > 0:
corners.copy_corners(qdel, "y", grid)
nx = grid.nic + 2 * nt # (grid.ie+nt) - (grid.is_-nt) + 1
ny = grid.njc + 2 * nt + 1 # (grid.je+nt+1) - (grid.js-nt) + 1
compute_meridional_flux(fy,
qdel,
grid.del6_u,
origin=origin,
domain=(nx, ny, km))
# Update q values
ny = grid.njc + 2 * nt # (grid.je+nt) - (grid.js-nt) + 1
update_q(qdel,
grid.rarea,
fx,
fy,
cd,
origin=origin,
domain=(nx, ny, km))
def compute(dp2, tracers, im, km, nq, jslice):
# Same as above, but with multiple tracer fields
i1 = spec.grid.is_
js = jslice.start
jext = jslice.stop - jslice.start
tracer_list = [tracers[q] for q in utils.tracer_variables[0:nq]]
shape = tracer_list[0].shape
shape_ij = shape[0:2]
dm = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
dm_pos = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
upper_fix = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
lower_fix = utils.make_storage_from_shape(shape, origin=(0, 0, 0))
zfix = utils.make_storage_from_shape(shape_ij, dtype=np.int, origin=(0, 0))
sum0 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
sum1 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
# TODO: Implement dev_gfs_physics ifdef when we implement compiler defs.
for tracer in tracer_list:
fix_tracer(
tracer,
dp2,
dm,
dm_pos,
zfix,
upper_fix,
lower_fix,
sum0,
sum1,
origin=(i1, js, 0),
domain=(im, jext, km),
)
return tracer_list
def test_storage_is_cached(backend, shape, origin, init):
outputs = []
for _ in range(2):
outputs.append(
make_storage_from_shape(
shape, origin=origin, init=init, cache_key="test-cached"
)
)
assert outputs[0] is outputs[1]
def compute_al(q, dxa, iord, is1, ie3, jfirst, jlast, kstart=0, nk=None):
if nk is None:
nk = grid().npz - kstart
dimensions = q.shape
local_origin = (origin[0], origin[1], kstart)
al = utils.make_storage_from_shape(dimensions, local_origin)
domain_y = (1, dimensions[1], nk)
if iord < 8:
main_al(
q,
al,
origin=(is1, jfirst, kstart),
domain=(ie3 - is1 + 1, jlast - jfirst + 1, nk),
)
if not grid().nested and spec.namelist.grid_type < 3:
if grid().west_edge:
al_y_edge_0(q,
dxa,
al,
origin=(grid().is_ - 1, 0, kstart),
domain=domain_y)
al_y_edge_1(q,
dxa,
al,
origin=(grid().is_, 0, kstart),
domain=domain_y)
al_y_edge_2(q,
dxa,
al,
origin=(grid().is_ + 1, 0, kstart),
domain=domain_y)
if grid().east_edge:
al_y_edge_0(q,
dxa,
al,
origin=(grid().ie, 0, kstart),
domain=domain_y)
al_y_edge_1(q,
dxa,
al,
origin=(grid().ie + 1, 0, kstart),
domain=domain_y)
al_y_edge_2(q,
dxa,
al,
origin=(grid().ie + 2, 0, kstart),
domain=domain_y)
if iord < 0:
floor_cap(
al,
0.0,
origin=(grid().is_ - 1, jfirst, kstart),
domain=(grid().nic + 3, jlast - jfirst + 1, nk),
)
return al
def compute(pk3, delp, ptop, akap):
grid = spec.grid
pei = utils.make_storage_from_shape(pk3.shape, grid.default_origin())
edge_domain_x = (2, grid.njc, grid.npz + 1)
edge_pe(
pei,
delp,
pk3,
ptop,
akap,
origin=(grid.is_ - 2, grid.js, 0),
domain=edge_domain_x,
)
edge_pe(
pei,
delp,
pk3,
ptop,
akap,
origin=(grid.ie + 1, grid.js, 0),
domain=edge_domain_x,
)
pej = utils.make_storage_from_shape(pk3.shape, grid.default_origin())
edge_domain_y = (grid.nic + 4, 2, grid.npz + 1)
edge_pe(
pej,
delp,
pk3,
ptop,
akap,
origin=(grid.is_ - 2, grid.js - 2, 0),
domain=edge_domain_y,
)
edge_pe(
pej,
delp,
pk3,
ptop,
akap,
origin=(grid.is_ - 2, grid.je + 1, 0),
domain=edge_domain_y,
)
def lagrangian_contributions_stencil(
q1: FloatField,
pe1: FloatField,
pe2: FloatField,
q4_1: FloatField,
q4_2: FloatField,
q4_3: FloatField,
q4_4: FloatField,
dp1: FloatField,
i1: int,
i2: int,
kord: int,
jslice: Tuple[int, int, int],
origin: Tuple[int, int, int],
domain: Tuple[int, int, int],
):
# A stencil with a loop over k2:
km = spec.grid.npz
shape2d = pe2.shape[0:2]
q2_adds = utils.make_storage_from_shape(shape2d)
ptop = utils.make_storage_from_shape(shape2d)
pbot = utils.make_storage_from_shape(shape2d)
for k_eul in range(km):
ptop[:, :] = pe2[:, :, k_eul]
pbot[:, :] = pe2[:, :, k_eul + 1]
q2_adds[:] = 0.0
lagrangian_contributions(
pe1,
ptop,
pbot,
q4_1,
q4_2,
q4_3,
q4_4,
dp1,
q2_adds,
origin=origin,
domain=domain,
)
q1[i1:i2 + 1, jslice, k_eul] = q2_adds[i1:i2 + 1, jslice]
def compute(self, inputs):
self.make_storage_data_input_vars(inputs)
ke = utils.make_storage_from_shape(inputs["uc"].shape)
vort = utils.make_storage_from_shape(inputs["vc"].shape)
c_sw.update_vorticity_and_kinetic_energy(
ke=ke,
vort=vort,
sin_sg1=self.grid.sin_sg1,
cos_sg1=self.grid.cos_sg1,
sin_sg2=self.grid.sin_sg2,
cos_sg2=self.grid.cos_sg2,
sin_sg3=self.grid.sin_sg3,
cos_sg3=self.grid.cos_sg3,
sin_sg4=self.grid.sin_sg4,
cos_sg4=self.grid.cos_sg4,
**inputs,
origin=(self.grid.is_ - 1, self.grid.js - 1, 0),
domain=(self.grid.nic + 2, self.grid.njc + 2, self.grid.npz),
)
return self.slice_output(inputs, {"ke_c": ke, "vort_c": vort})
def __init__(
self,
stencil_factory: StencilFactory,
check_negative: bool,
hydrostatic: bool,
):
grid_indexing = stencil_factory.grid_indexing
shape_ij = grid_indexing.domain_full(add=(1, 1, 0))[:2]
self._sum1 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
self._sum2 = utils.make_storage_from_shape(shape_ij, origin=(0, 0))
if check_negative:
raise NotImplementedError(
"Unimplemented namelist value check_negative=True")
if hydrostatic:
self._d0_vap = constants.CP_VAP - constants.C_LIQ
raise NotImplementedError(
"Unimplemented namelist hydrostatic=True")
else:
self._d0_vap = constants.CV_VAP - constants.C_LIQ
self._lv00 = constants.HLV - self._d0_vap * constants.TICE
self._fix_neg_water = stencil_factory.from_origin_domain(
func=fix_neg_water,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._fillq = stencil_factory.from_origin_domain(
func=fillq,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._fix_water_vapor_down = stencil_factory.from_origin_domain(
func=fix_water_vapor_down,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._fix_neg_cloud = stencil_factory.from_origin_domain(
func=fix_neg_cloud,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
def get_column_namelist(config: DGridShallowWaterLagrangianDynamicsConfig, npz):
"""
Generate a dictionary of columns that specify how parameters (such as nord, damp)
used in several functions called by D_SW vary over the k-dimension.
In a near-future PR, the need for this will disappear as we refactor
individual modules to apply this parameter variation explicitly in the
stencils themselves. If it doesn't, we should compute it only in the init phase.
The unique set of all column parameters is specified by k_bounds. For each k range
as specified by (kstart, nk) this sets what several different parameters are.
It previously was a dictionary with the k value as the key, the value being another
dictionary of values, but this did not work when we removed the k loop from some
modules and instead wanted to push the whole column ingestion down a level.
"""
direct_namelist = ["ke_bg", "d_con", "nord"]
all_names = direct_namelist + [
"nord_v",
"nord_w",
"nord_t",
"damp_vt",
"damp_w",
"damp_t",
"d2_divg",
]
col = {}
for name in all_names:
col[name] = utils.make_storage_from_shape((npz + 1,), (0,))
for name in direct_namelist:
col[name][:] = getattr(config, name)
col["d2_divg"][:] = min(0.2, config.d2_bg)
col["nord_v"][:] = min(2, col["nord"][0])
col["nord_w"][:] = col["nord_v"][0]
col["nord_t"][:] = col["nord_v"][0]
if config.do_vort_damp:
col["damp_vt"][:] = config.vtdm4
else:
col["damp_vt"][:] = 0
col["damp_w"][:] = col["damp_vt"][0]
col["damp_t"][:] = col["damp_vt"][0]
if npz == 1 or config.n_sponge < 0:
col["d2_divg"][0] = config.d2_bg
else:
col["d2_divg"][0] = max(0.01, config.d2_bg, config.d2_bg_k1)
lowest_kvals(col, 0, config.do_vort_damp)
if config.d2_bg_k2 > 0.01:
col["d2_divg"][1] = max(config.d2_bg, config.d2_bg_k2)
lowest_kvals(col, 1, config.do_vort_damp)
if config.d2_bg_k2 > 0.05:
col["d2_divg"][2] = max(config.d2_bg, 0.2 * config.d2_bg_k2)
set_low_kvals(col, 2)
return col
def compute(self, inputs):
i1 = self.grid.global_to_local_x(inputs["i1"] - 1)
i2 = self.grid.global_to_local_x(inputs["i2"] - 1)
j1 = 0
j2 = 0
self.compute_func = profile.RemapProfile(
self.grid.stencil_factory, inputs["kord"], inputs["iv"], i1, i2, j1, j2
)
self.make_storage_data_input_vars(inputs)
if "qs" not in inputs:
inputs["qs"] = utils.make_storage_from_shape(self.maxshape[0:2])
else:
qs_field = utils.make_storage_from_shape(self.maxshape[0:2], origin=(0, 0))
qs_field[i1 : i2 + 1, j1 : j2 + 1] = inputs["qs"][
i1 : i2 + 1, j1 : j2 + 1, 0
]
inputs["qs"] = qs_field
del inputs["km"], inputs["iv"], inputs["kord"], inputs["i1"], inputs["i2"]
q4_1, q4_2, q4_3, q4_4 = self.compute_func(**inputs)
return self.slice_output(
inputs, {"q4_1": q4_1, "q4_2": q4_2, "q4_3": q4_3, "q4_4": q4_4}
)
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 __init__(self, stencil_factory: StencilFactory, area):
grid_indexing = stencil_factory.grid_indexing
self._area = area
largest_possible_shape = grid_indexing.domain_full(add=(1, 1, 1))
self._gz_x = gt4py_utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)),
)
self._gz_y = gt4py_utils.make_storage_from_shape(
largest_possible_shape,
grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)),
)
full_origin = grid_indexing.origin_full()
full_domain = grid_indexing.domain_full(add=(0, 0, 1))
self._double_copy_stencil = stencil_factory.from_origin_domain(
double_copy,
origin=full_origin,
domain=full_domain,
)
ax_offsets = axis_offsets(grid_indexing, full_origin, full_domain)
self._fill_corners_x_stencil = stencil_factory.from_origin_domain(
corners.fill_corners_2cells_x_stencil,
externals=ax_offsets,
origin=full_origin,
domain=full_domain,
)
self._fill_corners_y_stencil = stencil_factory.from_origin_domain(
corners.fill_corners_2cells_y_stencil,
externals=ax_offsets,
origin=full_origin,
domain=full_domain,
)
self._update_dz_c = stencil_factory.from_origin_domain(
update_dz_c,
origin=grid_indexing.origin_compute(add=(-1, -1, 0)),
domain=grid_indexing.domain_compute(add=(2, 2, 1)),
)
def compute_from_storage(self, inputs):
wsd_2d = utils.make_storage_from_shape(inputs["wsd"].shape[0:2])
wsd_2d[:, :] = inputs["wsd"][:, :, 0]
inputs["wsd"] = wsd_2d
inputs["q_cld"] = inputs["tracers"]["qcld"]
l_to_e_obj = LagrangianToEulerian(
spec.grid.stencil_factory,
spec.namelist.dynamical_core.remapping,
spec.grid.area_64,
inputs["nq"],
inputs["pfull"],
)
l_to_e_obj(**inputs)
inputs.pop("q_cld")
return inputs
