def make_grid_storage(self, pygrid):
shape = pygrid.domain_shape_full(add=(1, 1, 1))
for k in TranslateGrid.composite_grid_vars:
if k in self.data:
self.make_composite_var_storage(k, self.data[k], shape)
del self.data[k]
for k, axis in TranslateGrid.edge_var_axis.items():
if k in self.data:
self.data[k] = utils.make_storage_data(
self.data[k],
shape,
start=(0, 0, pygrid.halo),
axis=axis,
read_only=True,
)
for k, v in self.data.items():
if type(v) is np.ndarray:
# TODO: when grid initialization model exists, may want to use
# it to inform this
istart, jstart = pygrid.horizontal_starts_from_shape(v.shape)
logger.debug("Storage for Grid variable {}, {}, {}, {}".format(
k, istart, jstart, v.shape))
origin = (istart, jstart, 0)
self.data[k] = utils.make_storage_data(v,
shape,
origin=origin,
start=origin,
read_only=True)
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_parallel(self, inputs, communicator):
# ak, bk, pfull, and phis are numpy arrays at this point and
# must be converted into gt4py storages
for name in ("ak", "bk", "pfull", "phis"):
inputs[name] = utils.make_storage_data(
inputs[name], inputs[name].shape,
len(inputs[name].shape) * (0, ))
grid_data = spec.grid.grid_data
grid_data.ak = inputs["ak"]
grid_data.bk = inputs["bk"]
self._base.compute_func = dyn_core.AcousticDynamics(
communicator,
spec.grid.stencil_factory,
grid_data,
spec.grid.damping_coefficients,
spec.grid.grid_type,
spec.grid.nested,
spec.grid.stretched_grid,
spec.namelist.dynamical_core.acoustic_dynamics,
inputs["ak"],
inputs["bk"],
inputs["pfull"],
inputs["phis"],
)
return super().compute_parallel(inputs, communicator)
def make_storage_data(
self,
array: np.ndarray,
istart: int = 0,
jstart: int = 0,
kstart: int = 0,
dummy_axes: Optional[Tuple[int, int, int]] = None,
axis: int = 2,
names_4d: Optional[List[str]] = None,
) -> Dict[str, type(Field)]:
use_shape = list(self.maxshape)
if dummy_axes:
for axis in dummy_axes:
use_shape[axis] = 1
use_shape = tuple(use_shape)
start = (istart, jstart, kstart)
if names_4d:
return utils.make_storage_dict(
array,
use_shape,
start=start,
origin=start,
dummy=dummy_axes,
axis=axis,
names=names_4d,
)
else:
return utils.make_storage_data(
array,
use_shape,
start=start,
origin=start,
dummy=dummy_axes,
axis=axis,
)
def compute_parallel(self, inputs, communicator):
# ak, bk, and phis are numpy arrays at this point and
# must be converted into gt4py storages
for name in ("ak", "bk", "phis"):
inputs[name] = utils.make_storage_data(
inputs[name], inputs[name].shape, len(inputs[name].shape) * (0,)
)
inputs["comm"] = communicator
state = self.state_from_inputs(inputs)
self.dycore = fv_dynamics.DynamicalCore(
comm=communicator,
grid_data=spec.grid.grid_data,
stencil_factory=spec.grid.stencil_factory,
damping_coefficients=spec.grid.damping_coefficients,
config=spec.namelist.dynamical_core,
ak=state["atmosphere_hybrid_a_coordinate"],
bk=state["atmosphere_hybrid_b_coordinate"],
phis=state["surface_geopotential"],
)
self.dycore.step_dynamics(
state,
inputs["consv_te"],
inputs["do_adiabatic_init"],
inputs["bdt"],
inputs["ptop"],
inputs["n_split"],
inputs["ks"],
)
outputs = self.outputs_from_state(state)
for name, value in outputs.items():
outputs[name] = self.subset_output(name, value)
return outputs
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 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 compute(self, inputs):
self.make_storage_data_input_vars(inputs)
inputs["index"] = utils.make_storage_data(
np.arange(satadjust.QS_LENGTH), self.maxshape, origin=(0, 0, 0)
)
kwargs = {"origin": (0, 0, 0), "domain": self.maxshape}
self.compute_func(**inputs, **kwargs)
for k, v in inputs.items():
if v.shape == self.maxshape:
inputs[k] = np.squeeze(v)
return inputs
def compute(self, inputs):
self.make_storage_data_input_vars(inputs)
index = np.arange(satadjust.QS_LENGTH)
inputs["index"] = utils.make_storage_data(index,
self.maxshape,
origin=(0, 0, 0),
read_only=False)
self._compute_q_tables_stencil(**inputs)
utils.device_sync()
for k, v in inputs.items():
if v.shape == self.maxshape:
inputs[k] = np.squeeze(v)
return inputs
def make_grid_storage(self, pygrid):
shape = pygrid.domain_shape_buffer_1cell()
for k in TranslateGrid.composite_grid_vars:
if k in self.data:
self.make_composite_var_storage(k, self.data[k], shape)
del self.data[k]
for k, axis in TranslateGrid.edge_var_axis.items():
if k in self.data:
edge_offset = pygrid.local_to_global_indices(pygrid.isd, pygrid.jsd)[
axis
]
if axis == 0:
edge_offset = pygrid.global_isd
width = pygrid.subtile_width_x
else:
edge_offset = pygrid.global_jsd
width = pygrid.subtile_width_y
edge_slice = slice(int(edge_offset), int(edge_offset + width + 1))
self.data[k] = utils.make_storage_data(
self.data[k][edge_slice],
shape,
start=(0, 0, pygrid.halo),
axis=axis,
)
for k, v in self.data.items():
if type(v) is np.ndarray:
# TODO: when grid initialization model exists, may want to use
# it to inform this
istart, jstart = pygrid.horizontal_starts_from_shape(v.shape)
logger.debug(
"Storage for Grid variable {}, {}, {}, {}".format(
k, istart, jstart, v.shape
)
)
origin = (istart, jstart, 0)
self.data[k] = utils.make_storage_data(
v, shape, origin=origin, start=origin
)
def compute(pt, pkz, heat_source, delz, delp, cappa, n_con, bdt):
grid = spec.grid
delt_column = np.ones(delz.shape[2]) * abs(bdt * spec.namelist.delt_max)
delt_column[0] *= 0.1
delt_column[1] *= 0.5
delt = utils.make_storage_data(
delt_column, delz.shape, origin=grid.default_origin()
)
compute_pkz_tempadjust(
delp,
delz,
cappa,
heat_source,
delt,
pt,
pkz,
origin=grid.compute_origin(),
domain=(grid.nic, grid.njc, n_con),
)
def make_storage_data(
self,
array: np.ndarray,
istart: int = 0,
jstart: int = 0,
kstart: int = 0,
dummy_axes: Optional[Tuple[int, int, int]] = None,
axis: int = 2,
names_4d: Optional[List[str]] = None,
read_only: bool = False,
full_shape: bool = False,
) -> Dict[str, "Field"]:
use_shape = list(self.maxshape)
if dummy_axes:
for axis in dummy_axes:
use_shape[axis] = 1
elif not full_shape and len(
array.shape) < 3 and axis == len(array.shape) - 1:
use_shape[1] = 1
start = (istart, jstart, kstart)
if names_4d:
return utils.make_storage_dict(
array,
use_shape,
start=start,
origin=start,
dummy=dummy_axes,
axis=axis,
names=names_4d,
)
else:
return utils.make_storage_data(
array,
use_shape,
start=start,
origin=start,
dummy=dummy_axes,
axis=axis,
read_only=read_only,
)
def __init__(
self,
stencil_factory: StencilFactory,
grid_data: GridData,
damping_coefficients: DampingCoefficients,
column_namelist,
nested: bool,
stretched_grid: bool,
config: DGridShallowWaterLagrangianDynamicsConfig,
):
self._f0 = spec.grid.f0
self.grid = grid_data
self.grid_indexing = stencil_factory.grid_indexing
assert config.grid_type < 3, "ubke and vbke only implemented for grid_type < 3"
assert not config.inline_q, "inline_q not yet implemented"
assert (
config.d_ext <= 0
), "untested d_ext > 0. need to call a2b_ord2, not yet implemented"
assert (column_namelist["damp_vt"] > dcon_threshold).all()
# TODO: in theory, we should check if damp_vt > 1e-5 for each k-level and
# only compute delnflux for k-levels where this is true
assert (column_namelist["damp_w"] > dcon_threshold).all()
# TODO: in theory, we should check if damp_w > 1e-5 for each k-level and
# only compute delnflux for k-levels where this is true
# only compute for k-levels where this is true
self.hydrostatic = config.hydrostatic
self._tmp_heat_s = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._vort_x_delta = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._vort_y_delta = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_ke = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_vort = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._uc_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._vc_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_ut = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_vt = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_fx = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_fy = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_gx = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_gy = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_dw = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_wk = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_fx2 = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_fy2 = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._tmp_damp_3d = utils.make_storage_from_shape(
(1, 1, self.grid_indexing.domain[2])
)
self._advected_u = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._advected_v = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._ub_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._vb_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self._column_namelist = column_namelist
self.delnflux_nosg_w = DelnFluxNoSG(
stencil_factory,
damping_coefficients,
grid_data.rarea,
self._column_namelist["nord_w"],
)
self.delnflux_nosg_v = DelnFluxNoSG(
stencil_factory,
damping_coefficients,
grid_data.rarea,
self._column_namelist["nord_v"],
)
self.fvtp2d_dp = FiniteVolumeTransport(
stencil_factory=stencil_factory,
grid_data=grid_data,
damping_coefficients=damping_coefficients,
grid_type=config.grid_type,
hord=config.hord_dp,
nord=self._column_namelist["nord_v"],
damp_c=self._column_namelist["damp_vt"],
)
self.fvtp2d_dp_t = FiniteVolumeTransport(
stencil_factory=stencil_factory,
grid_data=grid_data,
damping_coefficients=damping_coefficients,
grid_type=config.grid_type,
hord=config.hord_dp,
nord=self._column_namelist["nord_t"],
damp_c=self._column_namelist["damp_t"],
)
self.fvtp2d_tm = FiniteVolumeTransport(
stencil_factory=stencil_factory,
grid_data=grid_data,
damping_coefficients=damping_coefficients,
grid_type=config.grid_type,
hord=config.hord_tm,
nord=self._column_namelist["nord_v"],
damp_c=self._column_namelist["damp_vt"],
)
self.fvtp2d_vt_nodelnflux = FiniteVolumeTransport(
stencil_factory=stencil_factory,
grid_data=grid_data,
damping_coefficients=damping_coefficients,
grid_type=config.grid_type,
hord=config.hord_vt,
)
self.fv_prep = FiniteVolumeFluxPrep(
stencil_factory=stencil_factory,
grid_data=grid_data,
)
self.divergence_damping = DivergenceDamping(
stencil_factory,
grid_data,
damping_coefficients,
nested,
stretched_grid,
config.dddmp,
config.d4_bg,
config.nord,
config.grid_type,
column_namelist["nord"],
column_namelist["d2_divg"],
)
self._apply_pt_delp_fluxes = stencil_factory.from_dims_halo(
func=apply_pt_delp_fluxes_stencil_defn,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
externals={
"inline_q": config.inline_q,
},
)
self._kinetic_energy_update_part_1 = stencil_factory.from_dims_halo(
func=kinetic_energy_update_part_1,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
externals={
"iord": config.hord_mt,
"jord": config.hord_mt,
"mord": config.hord_mt,
"xt_minmax": False,
"yt_minmax": False,
},
skip_passes=("GreedyMerging",),
)
self._kinetic_energy_update_part_2 = stencil_factory.from_dims_halo(
func=kinetic_energy_update_part_2,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
externals={
"iord": config.hord_mt,
"jord": config.hord_mt,
"mord": config.hord_mt,
"xt_minmax": False,
"yt_minmax": False,
},
skip_passes=("GreedyMerging",),
)
self._flux_adjust_stencil = stencil_factory.from_dims_halo(
func=flux_adjust, compute_dims=[X_DIM, Y_DIM, Z_DIM]
)
self._flux_capacitor_stencil = stencil_factory.from_dims_halo(
func=flux_capacitor,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._vort_differencing_stencil = stencil_factory.from_dims_halo(
func=vort_differencing,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
)
self._u_and_v_from_ke_stencil = stencil_factory.from_dims_halo(
func=u_and_v_from_ke, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM]
)
self._compute_vort_stencil = stencil_factory.from_dims_halo(
func=compute_vort,
externals={
"radius": constants.RADIUS,
"do_f3d": config.do_f3d,
"hydrostatic": self.hydrostatic,
},
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._adjust_w_and_qcon_stencil = stencil_factory.from_dims_halo(
func=adjust_w_and_qcon,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
self._heat_diss_stencil = stencil_factory.from_dims_halo(
func=heat_diss,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
self._heat_source_from_vorticity_damping_stencil = (
stencil_factory.from_dims_halo(
func=heat_source_from_vorticity_damping,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
externals={
"do_skeb": config.do_skeb,
"d_con": config.d_con,
},
)
)
self._compute_vorticity_stencil = stencil_factory.from_dims_halo(
compute_vorticity,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._update_u_and_v_stencil = stencil_factory.from_dims_halo(
update_u_and_v, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM]
)
damping_factor_calculation_stencil = stencil_factory.from_origin_domain(
delnflux.calc_damp,
origin=(0, 0, 0),
domain=(1, 1, stencil_factory.grid_indexing.domain[2]),
)
damping_factor_calculation_stencil(
self._tmp_damp_3d,
self._column_namelist["nord_v"],
self._column_namelist["damp_vt"],
damping_coefficients.da_min_c,
)
self._delnflux_damp_vt = utils.make_storage_data(
self._tmp_damp_3d[0, 0, :], (self.grid_indexing.domain[2],), (0,)
)
damping_factor_calculation_stencil(
self._tmp_damp_3d,
self._column_namelist["nord_w"],
self._column_namelist["damp_w"],
damping_coefficients.da_min_c,
)
self._delnflux_damp_w = utils.make_storage_data(
self._tmp_damp_3d[0, 0, :], (self.grid_indexing.domain[2],), (0,)
)
y_temporary = utils.make_storage_from_shape(
shape=self.grid_indexing.max_shape, origin=self.grid_indexing.origin
)
self._copy_corners = PreAllocatedCopiedCornersFactory(
stencil_factory,
dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM],
y_temporary=y_temporary,
)
def __init__(
self,
comm: fv3gfs.util.CubedSphereCommunicator,
stencil_factory: StencilFactory,
grid_data: GridData,
damping_coefficients: DampingCoefficients,
grid_type,
nested,
stretched_grid,
config: AcousticDynamicsConfig,
# TODO: move ak, bk, pfull, and phis into GridData
ak: FloatFieldK,
bk: FloatFieldK,
pfull: FloatFieldK,
phis: FloatFieldIJ,
):
"""
Args:
comm: object for cubed sphere inter-process communication
stencil_factory: creates stencils
grid_data: metric terms defining the grid
damping_coefficients: damping configuration
grid_type: ???
nested: ???
stretched_grid: ???
config: configuration settings
ak: atmosphere hybrid a coordinate (Pa)
bk: atmosphere hybrid b coordinate (dimensionless)
pfull: atmospheric Eulerian grid reference pressure (Pa)
phis: surface geopotential height
"""
grid_indexing = stencil_factory.grid_indexing
self.comm = comm
self.config = config
assert config.d_ext == 0, "d_ext != 0 is not implemented"
assert config.beta == 0, "beta != 0 is not implemented"
assert not config.use_logp, "use_logp=True is not implemented"
self._da_min = damping_coefficients.da_min
self.grid_data = grid_data
self._pfull = pfull
self._nk_heat_dissipation = get_nk_heat_dissipation(
config.d_grid_shallow_water,
npz=grid_indexing.domain[2],
)
self.nonhydrostatic_pressure_gradient = (
nh_p_grad.NonHydrostaticPressureGradient(stencil_factory,
grid_data,
config.grid_type))
self._temporaries = dyncore_temporaries(grid_indexing)
self._temporaries["gz"][:] = HUGE_R
if not config.hydrostatic:
self._temporaries["pk3"][:] = HUGE_R
column_namelist = d_sw.get_column_namelist(config.d_grid_shallow_water,
grid_indexing.domain[2])
if not config.hydrostatic:
# To write lower dimensional storages, these need to be 3D
# then converted to lower dimensional
dp_ref_3d = utils.make_storage_from_shape(grid_indexing.max_shape)
zs_3d = utils.make_storage_from_shape(grid_indexing.max_shape)
dp_ref_stencil = stencil_factory.from_origin_domain(
dp_ref_compute,
origin=grid_indexing.origin_full(),
domain=grid_indexing.domain_full(add=(0, 0, 1)),
)
dp_ref_stencil(
ak,
bk,
phis,
dp_ref_3d,
zs_3d,
1.0 / constants.GRAV,
)
# After writing, make 'dp_ref' a K-field and 'zs' an IJ-field
self._dp_ref = utils.make_storage_data(dp_ref_3d[0, 0, :],
(dp_ref_3d.shape[2], ),
(0, ))
self._zs = utils.make_storage_data(zs_3d[:, :, 0],
zs_3d.shape[0:2], (0, 0))
self.update_height_on_d_grid = updatedzd.UpdateHeightOnDGrid(
stencil_factory,
damping_coefficients,
grid_data,
grid_type,
config.hord_tm,
self._dp_ref,
column_namelist,
d_sw.k_bounds(),
)
self.riem_solver3 = RiemannSolver3(stencil_factory, config.riemann)
self.riem_solver_c = RiemannSolverC(stencil_factory,
p_fac=config.p_fac)
origin, domain = grid_indexing.get_origin_domain(
[X_DIM, Y_DIM, Z_INTERFACE_DIM], halos=(2, 2))
self._compute_geopotential_stencil = stencil_factory.from_origin_domain(
compute_geopotential,
origin=origin,
domain=domain,
)
self.dgrid_shallow_water_lagrangian_dynamics = (
d_sw.DGridShallowWaterLagrangianDynamics(
stencil_factory,
grid_data,
damping_coefficients,
column_namelist,
nested,
stretched_grid,
config.d_grid_shallow_water,
))
self.cgrid_shallow_water_lagrangian_dynamics = CGridShallowWaterDynamics(
stencil_factory,
grid_data,
nested,
config.grid_type,
config.nord,
)
self._set_gz = stencil_factory.from_origin_domain(
set_gz,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(add=(0, 0, 1)),
)
self._set_pem = stencil_factory.from_origin_domain(
set_pem,
origin=grid_indexing.origin_compute(add=(-1, -1, 0)),
domain=grid_indexing.domain_compute(add=(2, 2, 0)),
)
self._p_grad_c = stencil_factory.from_origin_domain(
p_grad_c_stencil,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(add=(1, 1, 0)),
externals={"hydrostatic": config.hydrostatic},
)
self.update_geopotential_height_on_c_grid = (
updatedzc.UpdateGeopotentialHeightOnCGrid(stencil_factory,
grid_data.area))
self._zero_data = stencil_factory.from_origin_domain(
zero_data,
origin=grid_indexing.origin_full(),
domain=grid_indexing.domain_full(),
)
ax_offsets_pe = axis_offsets(
grid_indexing,
grid_indexing.origin_full(),
grid_indexing.domain_full(add=(0, 0, 1)),
)
self._edge_pe_stencil = stencil_factory.from_origin_domain(
pe_halo.edge_pe,
origin=grid_indexing.origin_full(),
domain=grid_indexing.domain_full(add=(0, 0, 1)),
externals={**ax_offsets_pe},
)
"""The stencil object responsible for updating the interface pressure"""
self._do_del2cubed = self._nk_heat_dissipation != 0 and config.d_con > 1.0e-5
if self._do_del2cubed:
nf_ke = min(3, config.nord + 1)
self._hyperdiffusion = HyperdiffusionDamping(stencil_factory,
damping_coefficients,
grid_data.rarea,
nmax=nf_ke)
if config.rf_fast:
self._rayleigh_damping = ray_fast.RayleighDamping(
stencil_factory,
rf_cutoff=config.rf_cutoff,
tau=config.tau,
hydrostatic=config.hydrostatic,
)
self._compute_pkz_tempadjust = stencil_factory.from_origin_domain(
temperature_adjust.compute_pkz_tempadjust,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.restrict_vertical(
nk=self._nk_heat_dissipation).domain_compute(),
)
self._pk3_halo = PK3Halo(stencil_factory)
self._copy_stencil = stencil_factory.from_origin_domain(
basic.copy_defn,
origin=grid_indexing.origin_full(),
domain=grid_indexing.domain_full(add=(0, 0, 1)),
)
# Halo updaters
self._halo_updaters = AcousticDynamics._HaloUpdaters(
self.comm, grid_indexing)
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 make_composite_var_storage(self, varname, data3d, shape):
for s in range(9):
self.data[varname + str(s + 1)] = utils.make_storage_data(
np.squeeze(data3d[:, :, s]), shape, origin=(0, 0, 0))
