def fvdyn_temporaries(shape):
grid = spec.grid
tmps = {}
halo_vars = ["cappa"]
storage_vars = ["te_2d", "dp1", "ph1", "ph2", "dp1", "wsd"]
column_vars = ["pfull", "gz", "cvm"]
plane_vars = ["te_2d", "te0_2d"]
utils.storage_dict(
tmps,
halo_vars + storage_vars + column_vars + plane_vars,
shape,
grid.default_origin(),
)
for q in halo_vars:
grid.quantity_dict_update(tmps, q)
return tmps
def dyncore_temporaries(shape):
grid = spec.grid
tmps = {}
utils.storage_dict(
tmps,
[
"ut", "vt", "gz", "zh", "pem", "ws3", "pkc", "pk3", "heat_source",
"divgd"
],
shape,
grid.default_origin(),
)
utils.storage_dict(tmps, ["crx", "xfx"], shape, grid.compute_x_origin())
utils.storage_dict(tmps, ["cry", "yfx"], shape, grid.compute_y_origin())
grid.quantity_dict_update(
tmps,
"heat_source",
dims=[fv3util.X_DIM, fv3util.Y_DIM, fv3util.Z_DIM])
for q in ["gz", "pkc", "zh"]:
grid.quantity_dict_update(
tmps,
q,
dims=[fv3util.X_DIM, fv3util.Y_DIM, fv3util.Z_INTERFACE_DIM])
grid.quantity_dict_update(
tmps,
"divgd",
dims=[fv3util.X_INTERFACE_DIM, fv3util.Y_INTERFACE_DIM, fv3util.Z_DIM],
)
return tmps
def dyncore_temporaries(grid_indexing: GridIndexing):
tmps: Dict[str, fv3gfs.util.Quantity] = {}
utils.storage_dict(
tmps,
["ut", "vt", "gz", "zh", "pem", "pkc", "pk3", "heat_source", "divgd"],
grid_indexing.max_shape,
grid_indexing.origin_full(),
)
utils.storage_dict(
tmps,
["ws3"],
grid_indexing.max_shape[0:2],
grid_indexing.origin_full()[0:2],
)
utils.storage_dict(
tmps,
["crx", "xfx"],
grid_indexing.max_shape,
grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)),
)
utils.storage_dict(
tmps,
["cry", "yfx"],
grid_indexing.max_shape,
grid_indexing.origin_compute(add=(-grid_indexing.n_halo, 0, 0)),
)
tmps["heat_source_quantity"] = quantity_wrap(tmps["heat_source"],
[X_DIM, Y_DIM, Z_DIM],
grid_indexing)
tmps["divgd_quantity"] = quantity_wrap(
tmps["divgd"],
dims=[fv3util.X_INTERFACE_DIM, fv3util.Y_INTERFACE_DIM, fv3util.Z_DIM],
grid_indexing=grid_indexing,
)
for name in ["gz", "pkc", "zh"]:
tmps[f"{name}_quantity"] = quantity_wrap(
tmps[name],
dims=[fv3util.X_DIM, fv3util.Y_DIM, fv3util.Z_INTERFACE_DIM],
grid_indexing=grid_indexing,
)
return tmps