def __init__(
self,
stencil_factory: StencilFactory,
grid_data: GridData,
):
grid_indexing = stencil_factory.grid_indexing
self._dx = grid_data.dx
self._dy = grid_data.dy
self._rdxa = grid_data.rdxa
self._rdya = grid_data.rdya
self._cosa_u = grid_data.cosa_u
self._cosa_v = grid_data.cosa_v
self._rsin_u = grid_data.rsin_u
self._rsin_v = grid_data.rsin_v
self._sin_sg1 = grid_data.sin_sg1
self._sin_sg2 = grid_data.sin_sg2
self._sin_sg3 = grid_data.sin_sg3
self._sin_sg4 = grid_data.sin_sg4
origin = grid_indexing.origin_full()
domain = grid_indexing.domain_full()
ax_offsets = axis_offsets(grid_indexing, origin, domain)
kwargs = {"externals": ax_offsets, "origin": origin, "domain": domain}
origin_corners = grid_indexing.origin_full(add=(1, 1, 0))
domain_corners = grid_indexing.domain_full(add=(-1, -1, 0))
corner_offsets = axis_offsets(grid_indexing, origin_corners,
domain_corners)
kwargs_corners = {
"externals": corner_offsets,
"origin": origin_corners,
"domain": domain_corners,
}
self._main_uc_contra_stencil = stencil_factory.from_origin_domain(
main_uc_contra, **kwargs)
self._main_vc_contra_stencil = stencil_factory.from_origin_domain(
main_vc_contra, **kwargs)
self._uc_contra_y_edge_stencil = stencil_factory.from_origin_domain(
uc_contra_y_edge, **kwargs)
self._vc_contra_y_edge_stencil = stencil_factory.from_origin_domain(
vc_contra_y_edge, **kwargs)
self._uc_contra_x_edge_stencil = stencil_factory.from_origin_domain(
uc_contra_x_edge, **kwargs)
self._vc_contra_x_edge_stencil = stencil_factory.from_origin_domain(
vc_contra_x_edge, **kwargs)
self._uc_contra_corners_stencil = stencil_factory.from_origin_domain(
uc_contra_corners, **kwargs_corners)
self._vc_contra_corners_stencil = stencil_factory.from_origin_domain(
vc_contra_corners, **kwargs_corners)
self._fxadv_fluxes_stencil = stencil_factory.from_origin_domain(
fxadv_fluxes_stencil, **kwargs)
def __init__(
self,
stencil_factory: StencilFactory,
grid_data: GridData,
grid_type: int,
jord: int,
):
if jord not in (5, 6, 7, 8):
raise NotImplementedError(
"Currently ytp_v is only supported for hord_mt == 5,6,7,8")
assert grid_type < 3
grid_indexing = stencil_factory.grid_indexing
origin = grid_indexing.origin_compute()
domain = grid_indexing.domain_compute(add=(1, 1, 0))
self._dy = grid_data.dy
self._dya = grid_data.dya
self._rdy = grid_data.rdy
ax_offsets = axis_offsets(grid_indexing, origin, domain)
self.stencil = stencil_factory.from_origin_domain(
ytp_v_stencil_defn,
externals={
"jord": jord,
"mord": jord,
"yt_minmax": False,
**ax_offsets,
},
origin=origin,
domain=domain,
skip_passes=("GreedyMerging", ),
)
def __init__(
self,
stencil_factory: StencilFactory,
dims: Sequence[str],
y_field,
) -> None:
"""
Args:
stencil_factory: creates stencils
dims: dimensionality of the data to be copied
y_field: 3D gt4py storage to use for y-differenceable field
(x-differenceable field uses same memory as base field)
"""
grid_indexing = stencil_factory.grid_indexing
origin, domain = grid_indexing.get_origin_domain(
dims=dims, halos=(grid_indexing.n_halo, grid_indexing.n_halo))
self._y_field = y_field
ax_offsets = axis_offsets(grid_indexing, origin, domain)
self._copy_corners_xy = stencil_factory.from_origin_domain(
func=copy_corners_xy_stencil_defn,
origin=origin,
domain=domain,
externals={
**ax_offsets,
},
)
def __init__(self, direction: str,
stencil_factory: StencilFactory) -> None:
"""The grid for this stencil"""
grid_indexing = stencil_factory.grid_indexing
n_halo = grid_indexing.n_halo
origin, domain = grid_indexing.get_origin_domain(
dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], halos=(n_halo, n_halo))
ax_offsets = axis_offsets(grid_indexing, origin, domain)
if direction == "x":
self._copy_corners = stencil_factory.from_origin_domain(
func=copy_corners_x_stencil_defn,
origin=origin,
domain=domain,
externals={
**ax_offsets,
},
)
elif direction == "y":
self._copy_corners = stencil_factory.from_origin_domain(
func=copy_corners_y_stencil_defn,
origin=origin,
domain=domain,
externals={
**ax_offsets,
},
)
else:
raise ValueError("Direction must be either 'x' or 'y'")
def __init__(self, grid):
super().__init__(grid)
self.in_vars["data_vars"] = {
"pe": {
"istart": grid.is_ - 1,
"iend": grid.ie + 1,
"jstart": grid.js - 1,
"jend": grid.je + 1,
"kend": grid.npz + 1,
"kaxis": 1,
},
"delp": {},
}
self.in_vars["parameters"] = ["ptop"]
self.out_vars = {"pe": self.in_vars["data_vars"]["pe"]}
ax_offsets_pe = axis_offsets(
self.grid.grid_indexing,
self.grid.grid_indexing.origin_full(),
self.grid.grid_indexing.domain_full(add=(0, 0, 1)),
)
self.compute_func = spec.grid.stencil_factory.from_origin_domain(
pe_halo.edge_pe,
origin=self.grid.grid_indexing.origin_full(),
domain=self.grid.grid_indexing.domain_full(add=(0, 0, 1)),
externals={**ax_offsets_pe},
)
def __init__(self, grid):
super().__init__(grid)
self.in_vars["data_vars"] = {
"uc": {},
"vc": {},
"ut": {},
"ub": grid.compute_dict_buffer_2d(),
}
self.in_vars["parameters"] = ["dt5", "dt4"]
self.out_vars = {"ub": grid.compute_dict_buffer_2d()}
origin = self.grid.compute_origin()
domain = self.grid.domain_shape_compute(add=(1, 1, 0))
ax_offsets = axis_offsets(self.grid, origin, domain)
self.compute_func = self.grid.stencil_factory.from_origin_domain(
ubke, externals=ax_offsets, origin=origin, domain=domain)
def __init__(self, stencil_factory: StencilFactory):
grid_indexing = stencil_factory.grid_indexing
origin = grid_indexing.origin_full()
domain = grid_indexing.domain_full(add=(0, 0, 1))
ax_offsets = axis_offsets(grid_indexing, origin, domain)
self._edge_pe_update = stencil_factory.from_origin_domain(
func=edge_pe_update,
externals={
**ax_offsets,
},
origin=origin,
domain=domain,
)
shape_2D = grid_indexing.domain_full(add=(1, 1, 1))[0:2]
self._pe_tmp = utils.make_storage_from_shape(
shape_2D, grid_indexing.origin_full())
def test_axis_offsets(
domain: Index3D,
n_halo: int,
south_edge: bool,
north_edge: bool,
west_edge: bool,
east_edge: bool,
origin_offset: Index3D,
domain_offset: Index3D,
i_start,
i_end,
j_start,
j_end,
):
grid = fv3core.utils.grid.GridIndexing(
domain=domain,
n_halo=n_halo,
south_edge=south_edge,
north_edge=north_edge,
west_edge=west_edge,
east_edge=east_edge,
)
origin = (n_halo, n_halo, 0)
call_origin = tuple(compute + offset
for (compute, offset) in zip(origin, origin_offset))
call_domain = tuple(compute + offset
for (compute, offset) in zip(domain, domain_offset))
axis_offsets = grid.axis_offsets(call_origin, call_domain)
if west_edge:
assert axis_offsets["i_start"] == i_start
else:
assert axis_offsets["i_start"] == gtscript.I[0] - np.iinfo(
np.int16).max
if east_edge:
assert axis_offsets["i_end"] == i_end
else:
assert axis_offsets["i_end"] == gtscript.I[-1] + np.iinfo(np.int16).max
if south_edge:
assert axis_offsets["j_start"] == j_start
else:
assert axis_offsets["j_start"] == gtscript.J[0] - np.iinfo(
np.int16).max
if north_edge:
assert axis_offsets["j_end"] == j_end
else:
assert axis_offsets["j_end"] == gtscript.J[-1] + np.iinfo(np.int16).max
def __init__(
self,
direction: str,
stencil_factory: StencilFactory,
temporary_field=None,
origin=None,
domain=None,
) -> None:
n_halo = stencil_factory.grid_indexing.n_halo
(
default_origin,
default_domain,
) = stencil_factory.grid_indexing.get_origin_domain(
dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_INTERFACE_DIM],
halos=(n_halo, n_halo),
)
if origin is None:
origin = default_origin
"""The origin for the corner computation"""
if domain is None:
domain = default_domain
"""The full domain required to do corner computation everywhere"""
if temporary_field is not None:
self._corner_tmp = temporary_field
else:
self._corner_tmp = utils.make_storage_from_shape(
stencil_factory.grid_indexing.max_shape, origin=origin)
if direction == "x":
defn = fill_corners_bgrid_x_defn
elif direction == "y":
defn = fill_corners_bgrid_y_defn
else:
raise ValueError("Direction must be either 'x' or 'y'")
externals = axis_offsets(stencil_factory.grid_indexing,
origin=origin,
domain=domain)
self._fill_corners_bgrid = stencil_factory.from_origin_domain(
func=defn, origin=origin, domain=domain, externals=externals)
def compute(self, inputs):
self.make_storage_data_input_vars(inputs)
origin = self.grid.full_origin()
domain = self.grid.domain_shape_full()
axes_offsets = axis_offsets(self.grid, origin, domain)
if inputs["dir"] == 1:
stencil = self.grid.stencil_factory.from_origin_domain(
corners.fill_corners_2cells_x_stencil,
externals=axes_offsets,
origin=origin,
domain=domain,
)
elif inputs["dir"] == 2:
stencil = self.grid.stencil_factory.from_origin_domain(
corners.fill_corners_2cells_y_stencil,
externals=axes_offsets,
origin=origin,
domain=domain,
)
stencil(inputs["q4c"], inputs["q4c"])
return self.slice_output(inputs, {"q4c": inputs["q4c"]})
def compute(self, inputs):
nord_column = inputs["nord_col"][:]
self.make_storage_data_input_vars(inputs)
for nord in utils.unique(nord_column):
if nord != 0:
ki = [
k for k in range(self.grid.npz)
if nord_column[0, 0, k] == nord
]
origin = (self.grid.isd, self.grid.jsd, ki[0])
domain = (self.grid.nid + 1, self.grid.njd + 1, len(ki))
axes_offsets = axis_offsets(self.grid.grid_indexing, origin,
domain)
vector_corner_fill = self.grid.stencil_factory.from_origin_domain(
corners.fill_corners_dgrid_defn,
externals=axes_offsets,
origin=origin,
domain=domain,
)
vector_corner_fill(inputs["vc"], inputs["vc"], inputs["uc"],
inputs["uc"], -1.0)
return self.slice_output(inputs)
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 __init__(
self,
stencil_factory: StencilFactory,
damping_coefficients: DampingCoefficients,
rarea,
nord,
nk: Optional[int] = None,
):
"""
nord sets the order of damping to apply:
nord = 0: del-2
nord = 1: del-4
nord = 2: del-6
"""
grid_indexing = stencil_factory.grid_indexing
self._del6_u = damping_coefficients.del6_u
self._del6_v = damping_coefficients.del6_v
self._rarea = rarea
if max(nord[:]) > 3:
raise ValueError("nord must be less than 3")
if not np.all(n in [0, 2, 3] for n in nord[:]):
raise NotImplementedError("nord must have values 0, 2, or 3")
self._nmax = int(max(nord[:]))
i1 = grid_indexing.isc - 1 - self._nmax
i2 = grid_indexing.iec + 1 + self._nmax
j1 = grid_indexing.jsc - 1 - self._nmax
j2 = grid_indexing.jec + 1 + self._nmax
if nk is None:
nk = grid_indexing.domain[2]
nk = nk
origin_d2 = (i1, j1, 0)
domain_d2 = (i2 - i1 + 1, j2 - j1 + 1, nk)
f1_ny = grid_indexing.jec - grid_indexing.jsc + 1 + 2 * self._nmax
f1_nx = grid_indexing.iec - grid_indexing.isc + 2 + 2 * self._nmax
fx_origin = (grid_indexing.isc - self._nmax, grid_indexing.jsc - self._nmax, 0)
self._nord = nord
if nk <= 3:
raise NotImplementedError("nk must be more than 3 for DelnFluxNoSG")
self._k_bounds = [1, 1, 1, nk - 3]
preamble_ax_offsets = axis_offsets(grid_indexing, origin_d2, domain_d2)
fx_ax_offsets = axis_offsets(grid_indexing, fx_origin, (f1_nx, f1_ny, nk))
fy_ax_offsets = axis_offsets(
grid_indexing, fx_origin, (f1_nx - 1, f1_ny + 1, nk)
)
origins_d2 = []
domains_d2 = []
origins_flux = []
domains_fx = []
domains_fy = []
for n in range(self._nmax):
nt = self._nmax - 1 - n
nt_ny = grid_indexing.jec - grid_indexing.jsc + 3 + 2 * nt
nt_nx = grid_indexing.iec - grid_indexing.isc + 3 + 2 * nt
origins_d2.append(
(grid_indexing.isc - nt - 1, grid_indexing.jsc - nt - 1, 0)
)
domains_d2.append((nt_nx, nt_ny, nk))
origins_flux.append((grid_indexing.isc - nt, grid_indexing.jsc - nt, 0))
domains_fx.append((nt_nx - 1, nt_ny - 2, nk))
domains_fy.append((nt_nx - 2, nt_ny - 1, nk))
nord_dictionary = {
"nord0": nord[0],
"nord1": nord[1],
"nord2": nord[2],
"nord3": nord[3],
}
self._d2_damp = stencil_factory.from_origin_domain(
d2_damp_interval,
externals={
**nord_dictionary,
**preamble_ax_offsets,
},
origin=origin_d2,
domain=domain_d2,
)
self._copy_stencil_interval = stencil_factory.from_origin_domain(
copy_stencil_interval,
externals={
**nord_dictionary,
**preamble_ax_offsets,
},
origin=origin_d2,
domain=domain_d2,
)
self._d2_stencil = get_stencils_with_varied_bounds(
d2_highorder_stencil,
origins_d2,
domains_d2,
stencil_factory=stencil_factory,
externals={**nord_dictionary},
)
self._column_conditional_fx_calculation = get_stencils_with_varied_bounds(
fx_calc_stencil_column,
origins_flux,
domains_fx,
stencil_factory=stencil_factory,
externals={**nord_dictionary},
)
self._column_conditional_fy_calculation = get_stencils_with_varied_bounds(
fy_calc_stencil_column,
origins_flux,
domains_fy,
stencil_factory=stencil_factory,
externals={**nord_dictionary},
)
self._fx_calc_stencil = stencil_factory.from_origin_domain(
fx_calc_stencil_nord,
externals={**fx_ax_offsets, **nord_dictionary},
origin=fx_origin,
domain=(f1_nx, f1_ny, nk),
)
self._fy_calc_stencil = stencil_factory.from_origin_domain(
fy_calc_stencil_nord,
externals={**fy_ax_offsets, **nord_dictionary},
origin=fx_origin,
domain=(f1_nx - 1, f1_ny + 1, nk),
)
corner_origin, corner_domain = grid_indexing.get_origin_domain(
dims=[X_DIM, Y_DIM, Z_DIM],
halos=(grid_indexing.n_halo, grid_indexing.n_halo),
)
corner_domain = corner_domain[:2] + (nk,)
corner_axis_offsets = axis_offsets(grid_indexing, corner_origin, corner_domain)
self._corner_tmp = utils.make_storage_from_shape(
corner_domain, origin=corner_origin
)
self._copy_corners_x_nord = stencil_factory.from_origin_domain(
copy_corners_x_nord,
externals={**corner_axis_offsets, **nord_dictionary},
origin=corner_origin,
domain=corner_domain,
)
self._copy_corners_y_nord = stencil_factory.from_origin_domain(
copy_corners_y_nord,
externals={**corner_axis_offsets, **nord_dictionary},
origin=corner_origin,
domain=corner_domain,
)
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 __init__(
self,
stencil_factory: StencilFactory,
config: RemappingConfig,
area_64,
nq,
pfull,
):
grid_indexing = stencil_factory.grid_indexing
if config.kord_tm >= 0:
raise NotImplementedError(
"map ppm, untested mode where kord_tm >= 0")
hydrostatic = config.hydrostatic
if hydrostatic:
raise NotImplementedError("Hydrostatic is not implemented")
shape_kplus = grid_indexing.domain_full(add=(0, 0, 1))
self._t_min = 184.0
self._nq = nq
# do_omega = hydrostatic and last_step # TODO pull into inputs
self._domain_jextra = (
grid_indexing.domain[0],
grid_indexing.domain[1] + 1,
grid_indexing.domain[2] + 1,
)
self._pe1 = utils.make_storage_from_shape(shape_kplus)
self._pe2 = utils.make_storage_from_shape(shape_kplus)
self._dp2 = utils.make_storage_from_shape(shape_kplus)
self._pn2 = utils.make_storage_from_shape(shape_kplus)
self._pe0 = utils.make_storage_from_shape(shape_kplus)
self._pe3 = utils.make_storage_from_shape(shape_kplus)
self._gz: FloatField = utils.make_storage_from_shape(
shape_kplus, grid_indexing.origin_compute())
self._cvm: FloatField = utils.make_storage_from_shape(
shape_kplus, grid_indexing.origin_compute())
self._init_pe = stencil_factory.from_origin_domain(
init_pe,
origin=grid_indexing.origin_compute(),
domain=self._domain_jextra)
self._moist_cv_pt_pressure = stencil_factory.from_origin_domain(
moist_cv_pt_pressure,
externals={
"kord_tm": config.kord_tm,
"hydrostatic": hydrostatic
},
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(add=(0, 0, 1)),
)
self._moist_cv_pkz = stencil_factory.from_origin_domain(
moist_cv.moist_pkz,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._pn2_pk_delp = stencil_factory.from_origin_domain(
pn2_pk_delp,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._kord_tm = abs(config.kord_tm)
self._map_single_pt = MapSingle(
stencil_factory,
self._kord_tm,
1,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec,
)
self._mapn_tracer = MapNTracer(
stencil_factory,
abs(config.kord_tr),
nq,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec,
fill=config.fill,
)
self._kord_wz = config.kord_wz
self._map_single_w = MapSingle(
stencil_factory,
self._kord_wz,
-2,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec,
)
self._map_single_delz = MapSingle(
stencil_factory,
self._kord_wz,
1,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec,
)
self._undo_delz_adjust_and_copy_peln = stencil_factory.from_origin_domain(
undo_delz_adjust_and_copy_peln,
origin=grid_indexing.origin_compute(),
domain=(
grid_indexing.domain[0],
grid_indexing.domain[1],
grid_indexing.domain[2] + 1,
),
)
self._pressures_mapu = stencil_factory.from_origin_domain(
pressures_mapu,
origin=grid_indexing.origin_compute(),
domain=self._domain_jextra,
)
self._kord_mt = config.kord_mt
self._map_single_u = MapSingle(
stencil_factory,
self._kord_mt,
-1,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec + 1,
)
self._pressures_mapv = stencil_factory.from_origin_domain(
pressures_mapv,
origin=grid_indexing.origin_compute(),
domain=(
grid_indexing.domain[0] + 1,
grid_indexing.domain[1],
grid_indexing.domain[2] + 1,
),
)
self._map_single_v = MapSingle(
stencil_factory,
self._kord_mt,
-1,
grid_indexing.isc,
grid_indexing.iec + 1,
grid_indexing.jsc,
grid_indexing.jec,
)
ax_offsets_jextra = axis_offsets(
grid_indexing,
grid_indexing.origin_compute(),
grid_indexing.domain_compute(add=(0, 1, 0)),
)
self._update_ua = stencil_factory.from_origin_domain(
update_ua,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(add=(0, 1, 0)),
externals={**ax_offsets_jextra},
)
self._copy_from_below_stencil = stencil_factory.from_origin_domain(
copy_from_below,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._moist_cv_last_step_stencil = stencil_factory.from_origin_domain(
moist_pt_last_step,
origin=(grid_indexing.isc, grid_indexing.jsc, 0),
domain=(
grid_indexing.domain[0],
grid_indexing.domain[1],
grid_indexing.domain[2] + 1,
),
)
self._basic_adjust_divide_stencil = stencil_factory.from_origin_domain(
adjust_divide_stencil,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(),
)
self._do_sat_adjust = config.do_sat_adj
self.kmp = grid_indexing.domain[2] - 1
for k in range(pfull.shape[0]):
if pfull[k] > 10.0e2:
self.kmp = k
break
self._saturation_adjustment = SatAdjust3d(stencil_factory,
config.sat_adjust, area_64,
self.kmp)
self._sum_te_stencil = stencil_factory.from_origin_domain(
sum_te,
origin=(grid_indexing.isc, grid_indexing.jsc, self.kmp),
domain=(
grid_indexing.domain[0],
grid_indexing.domain[1],
grid_indexing.domain[2] - self.kmp,
),
)
def __init__(
self,
stencil_factory: StencilFactory,
grid_data: GridData,
nested: bool,
grid_type: int,
dord4: bool,
):
grid_indexing = stencil_factory.grid_indexing
self._cosa_s = grid_data.cosa_s
self._cosa_u = grid_data.cosa_u
self._cosa_v = grid_data.cosa_v
self._rsin_u = grid_data.rsin_u
self._rsin_v = grid_data.rsin_v
self._rsin2 = grid_data.rsin2
self._dxa = grid_data.dxa
self._dya = grid_data.dya
self._sin_sg1 = grid_data.sin_sg1
self._sin_sg2 = grid_data.sin_sg2
self._sin_sg3 = grid_data.sin_sg3
self._sin_sg4 = grid_data.sin_sg4
if grid_type >= 3:
raise NotImplementedError("unimplemented grid_type >= 3")
self._big_number = 1e30 # 1e8 if 32 bit
nx = grid_indexing.iec + 1 # grid.npx + 2
ny = grid_indexing.jec + 1 # grid.npy + 2
i1 = grid_indexing.isc - 1
j1 = grid_indexing.jsc - 1
id_ = 1 if dord4 else 0
pad = 2 + 2 * id_
npt = 4 if not nested else 0
if npt > grid_indexing.domain[0] - 1 or npt > grid_indexing.domain[
1] - 1:
npt = 0
self._utmp = utils.make_storage_from_shape(
grid_indexing.max_shape,
grid_indexing.origin_full(),
)
self._vtmp = utils.make_storage_from_shape(grid_indexing.max_shape,
grid_indexing.origin_full())
js1 = npt + OFFSET if grid_indexing.south_edge else grid_indexing.jsc - 1
je1 = ny - npt if grid_indexing.north_edge else grid_indexing.jec + 1
is1 = npt + OFFSET if grid_indexing.west_edge else grid_indexing.isd
ie1 = nx - npt if grid_indexing.east_edge else grid_indexing.ied
is2 = npt + OFFSET if grid_indexing.west_edge else grid_indexing.isc - 1
ie2 = nx - npt if grid_indexing.east_edge else grid_indexing.iec + 1
js2 = npt + OFFSET if grid_indexing.south_edge else grid_indexing.jsd
je2 = ny - npt if grid_indexing.north_edge else grid_indexing.jed
ifirst = (grid_indexing.isc +
2 if grid_indexing.west_edge else grid_indexing.isc - 1)
ilast = (grid_indexing.iec -
1 if grid_indexing.east_edge else grid_indexing.iec + 2)
idiff = ilast - ifirst + 1
jfirst = (grid_indexing.jsc +
2 if grid_indexing.south_edge else grid_indexing.jsc - 1)
jlast = (grid_indexing.jec -
1 if grid_indexing.north_edge else grid_indexing.jec + 2)
jdiff = jlast - jfirst + 1
self._set_tmps = stencil_factory.from_dims_halo(
func=set_tmps,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._lagrange_interpolation_y_p1 = stencil_factory.from_origin_domain(
func=lagrange_interpolation_y_p1,
origin=(is1, js1, 0),
domain=(ie1 - is1 + 1, je1 - js1 + 1, grid_indexing.domain[2]),
)
self._lagrange_interpolation_x_p1 = stencil_factory.from_origin_domain(
func=lagrange_interpolation_x_p1,
origin=(is2, js2, 0),
domain=(ie2 - is2 + 1, je2 - js2 + 1, grid_indexing.domain[2]),
)
origin = grid_indexing.origin_full()
domain = grid_indexing.domain_full()
ax_offsets = axis_offsets(grid_indexing, origin, domain)
if npt == 0:
d2a2c_avg_offset = -1
else:
d2a2c_avg_offset = 3
self._avg_box = stencil_factory.from_dims_halo(
func=avg_box,
externals={"D2A2C_AVG_OFFSET": d2a2c_avg_offset},
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._contravariant_components = stencil_factory.from_origin_domain(
func=contravariant_components,
origin=grid_indexing.origin_compute(add=(-1 - id_, -1 - id_, 0)),
domain=grid_indexing.domain_compute(add=(pad, pad, 0)),
)
origin_edges = grid_indexing.origin_compute(add=(-3, -3, 0))
domain_edges = grid_indexing.domain_compute(add=(6, 6, 0))
ax_offsets_edges = axis_offsets(grid_indexing, origin_edges,
domain_edges)
self._fill_corners_x = stencil_factory.from_origin_domain(
func=fill_corners_x,
externals=ax_offsets_edges,
origin=origin_edges,
domain=domain_edges,
)
self._ut_main = stencil_factory.from_origin_domain(
func=ut_main,
origin=(ifirst, j1, 0),
domain=(idiff, grid_indexing.domain[1] + 2,
grid_indexing.domain[2]),
)
self._east_west_edges = stencil_factory.from_origin_domain(
func=east_west_edges,
externals={
"i_end": ax_offsets_edges["i_end"],
"i_start": ax_offsets_edges["i_start"],
"local_je": ax_offsets_edges["local_je"],
"local_js": ax_offsets_edges["local_js"],
},
origin=origin_edges,
domain=domain_edges,
)
# Ydir:
self._fill_corners_y = stencil_factory.from_origin_domain(
func=fill_corners_y,
externals={
**ax_offsets_edges,
},
origin=origin_edges,
domain=domain_edges,
)
self._north_south_edges = stencil_factory.from_origin_domain(
func=north_south_edges,
externals={
"j_end": ax_offsets_edges["j_end"],
"j_start": ax_offsets_edges["j_start"],
"local_ie": ax_offsets_edges["local_ie"],
"local_is": ax_offsets_edges["local_is"],
"local_je": ax_offsets_edges["local_je"],
"local_js": ax_offsets_edges["local_js"],
},
origin=origin_edges,
domain=domain_edges,
)
self._vt_main = stencil_factory.from_origin_domain(
func=vt_main,
origin=(i1, jfirst, 0),
domain=(grid_indexing.domain[0] + 2, jdiff,
grid_indexing.domain[2]),
)
