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,
stencil_factory: StencilFactory,
damping_coefficients: DampingCoefficients,
grid_data: GridData,
grid_type: int,
hord_tm: int,
dp0: FloatFieldK,
column_namelist,
k_bounds,
):
"""
Args:
grid: fv3core grid object
namelist: flattened fv3gfs namelist
dp0: air pressure on interface levels, reference pressure
can be used as an approximation
column_namelist: ???
k_bounds: ???
"""
grid_indexing = stencil_factory.grid_indexing
self.grid_indexing = grid_indexing
self._area = grid_data.area
self._column_namelist = column_namelist
self._k_bounds = k_bounds # d_sw.k_bounds()
if any(column_namelist["damp_vt"][kstart] <= 1e-5
for kstart in range(len(k_bounds))):
raise NotImplementedError(
"damp <= 1e-5 in column_cols is untested")
self._dp0 = dp0
self._allocate_temporary_storages(grid_indexing)
self._initialize_interpolation_constants(
stencil_factory=stencil_factory, grid_indexing=grid_indexing)
self._interpolate_to_layer_interface = stencil_factory.from_origin_domain(
cubic_spline_interpolation_from_layer_center_to_interfaces,
origin=grid_indexing.origin_full(),
domain=grid_indexing.domain_full(add=(0, 0, 1)),
)
self._apply_height_fluxes = stencil_factory.from_origin_domain(
apply_height_fluxes,
origin=grid_indexing.origin_compute(),
domain=grid_indexing.domain_compute(add=(0, 0, 1)),
)
self.delnflux = DelnFluxNoSG(
stencil_factory,
damping_coefficients,
grid_data.rarea,
self._column_namelist["nord_v"],
nk=grid_indexing.domain[2] + 1,
)
self.finite_volume_transport = FiniteVolumeTransport(
stencil_factory=stencil_factory,
grid_data=grid_data,
damping_coefficients=damping_coefficients,
grid_type=grid_type,
hord=hord_tm,
)
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 __init__(
self,
stencil_factory: StencilFactory,
kord: int,
mode: int,
i1: int,
i2: int,
j1: int,
j2: int,
):
# TODO: consider refactoring to take in origin and domain
grid_indexing = stencil_factory.grid_indexing
shape = grid_indexing.domain_full(add=(1, 1, 1))
origin = grid_indexing.origin_compute()
self._dp1 = utils.make_storage_from_shape(shape, origin=origin)
self._q4_1 = utils.make_storage_from_shape(shape, origin=origin)
self._q4_2 = utils.make_storage_from_shape(shape, origin=origin)
self._q4_3 = utils.make_storage_from_shape(shape, origin=origin)
self._q4_4 = utils.make_storage_from_shape(shape, origin=origin)
self._tmp_qs = utils.make_storage_from_shape(shape[0:2], origin=(0, 0))
self._lev = utils.make_storage_from_shape(
shape[:-1],
origin=origin[:-1],
mask=(True, True, False),
dtype=int,
)
self._extents = (i2 - i1 + 1, j2 - j1 + 1)
origin = (i1, j1, 0)
domain = (*self._extents, grid_indexing.domain[2])
self._lagrangian_contributions = stencil_factory.from_origin_domain(
lagrangian_contributions,
origin=origin,
domain=domain,
)
self._remap_profile = RemapProfile(stencil_factory, kord, mode, i1, i2,
j1, j2)
self._set_dp = stencil_factory.from_origin_domain(set_dp,
origin=origin,
domain=domain)
self._copy_stencil = stencil_factory.from_origin_domain(
copy_defn,
origin=(0, 0, 0),
domain=grid_indexing.domain_full(),
)
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 get_stencils_with_varied_bounds(
func: Callable[..., None],
origins: List[Index3D],
domains: List[Index3D],
stencil_factory: StencilFactory,
externals: Optional[Mapping[str, Any]] = None,
) -> List[FrozenStencil]:
assert len(origins) == len(domains), (
"Lists of origins and domains need to have the same length, you provided "
+ str(len(origins))
+ " origins and "
+ str(len(domains))
+ " domains"
)
if externals is None:
externals = {}
stencils = []
for origin, domain in zip(origins, domains):
ax_offsets = fv3core.utils.grid.axis_offsets(spec.grid, origin, domain)
stencils.append(
stencil_factory.from_origin_domain(
func,
origin=origin,
domain=domain,
externals={**externals, **ax_offsets},
)
)
return stencils
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,
stencil_factory: StencilFactory,
dxa,
grid_type: int,
iord,
origin: Index3D,
domain: Index3D,
):
# Arguments come from:
# namelist.grid_type
# grid.dxa
assert grid_type < 3
self._dxa = dxa
ax_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain)
self._compute_flux_stencil = stencil_factory.from_origin_domain(
func=compute_x_flux,
externals={
"iord": iord,
"mord": abs(iord),
"xt_minmax": True,
"i_start": ax_offsets["i_start"],
"i_end": ax_offsets["i_end"],
},
origin=origin,
domain=domain,
)
def __init__(
self, stencil_factory: StencilFactory, grid_data: GridData, order: int
):
"""
Initializes stencils to use either 2nd or 4th order of interpolation
based on namelist setting
Args:
stencil_factory: creates stencils
grid_data: object with metric terms
order: Order of interpolation, must be 2 or 4
"""
grid_indexing = stencil_factory.grid_indexing
self._n_halo = grid_indexing.n_halo
self._dx = grid_data.dx
self._dy = grid_data.dy
# TODO: define these based on data from grid_data
self._a11 = spec.grid.a11
self._a12 = spec.grid.a12
self._a21 = spec.grid.a21
self._a22 = spec.grid.a22
if order == 2:
self._do_ord4 = False
halos = (1, 1)
func = c2l_ord2
else:
self._do_ord4 = True
halos = (0, 0)
func = ord4_transform
self._compute_cubed_to_latlon = stencil_factory.from_dims_halo(
func=func, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=halos
)
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 _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 __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, p_fac, istart, iend,
jstart, jend, nk):
self._pfac = p_fac
nic = iend - istart + 1
njc = jend - jstart + 1
self._compute_sim1_solve = stencil_factory.from_origin_domain(
func=sim1_solver,
origin=(istart, jstart, 0),
domain=(nic, njc, nk),
)
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 __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, config: RiemannConfig):
grid_indexing = stencil_factory.grid_indexing
self._sim1_solve = Sim1Solver(
stencil_factory,
config.p_fac,
grid_indexing.isc,
grid_indexing.iec,
grid_indexing.jsc,
grid_indexing.jec,
grid_indexing.domain[2] + 1,
)
if config.a_imp <= 0.999:
raise NotImplementedError("a_imp <= 0.999 is not implemented")
riemorigin = grid_indexing.origin_compute()
domain = grid_indexing.domain_compute(add=(0, 0, 1))
shape = grid_indexing.max_shape
self._tmp_dm = utils.make_storage_from_shape(shape, riemorigin)
self._tmp_pe_init = utils.make_storage_from_shape(shape, riemorigin)
self._tmp_pm = utils.make_storage_from_shape(shape, riemorigin)
self._tmp_pem = utils.make_storage_from_shape(shape, riemorigin)
self._tmp_peln_run = utils.make_storage_from_shape(shape, riemorigin)
self._tmp_gm = utils.make_storage_from_shape(shape, riemorigin)
self._precompute_stencil = stencil_factory.from_origin_domain(
precompute,
origin=riemorigin,
domain=domain,
)
self._finalize_stencil = stencil_factory.from_origin_domain(
finalize,
externals={
"use_logp": config.use_logp,
"beta": config.beta
},
origin=riemorigin,
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 get_stencil_factory(backend) -> StencilFactory:
config = StencilConfig(
backend=backend,
rebuild=False,
validate_args=False,
format_source=False,
device_sync=False,
)
indexing = GridIndexing(
domain=(12, 12, 79),
n_halo=3,
south_edge=False,
north_edge=False,
west_edge=False,
east_edge=False,
)
return StencilFactory(config=config, grid_indexing=indexing)
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 __init__(self, stencil_factory: StencilFactory, rf_cutoff, tau,
hydrostatic):
grid_indexing = stencil_factory.grid_indexing
self._rf_cutoff = rf_cutoff
origin, domain = grid_indexing.get_origin_domain(
[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM])
ax_offsets = axis_offsets(grid_indexing, origin, domain)
local_axis_offsets = {}
for axis_offset_name, axis_offset_value in ax_offsets.items():
if "local" in axis_offset_name:
local_axis_offsets[axis_offset_name] = axis_offset_value
self._ray_fast_wind_compute = stencil_factory.from_origin_domain(
ray_fast_wind_compute,
origin=origin,
domain=domain,
externals={
"hydrostatic": hydrostatic,
"rf_cutoff": rf_cutoff,
"tau": tau,
**local_axis_offsets,
},
)
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,
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,
stencil_factory: StencilFactory,
nwat: int,
fv_sg_adj: float,
n_sponge: int,
hydrostatic: bool,
):
assert not hydrostatic, "Hydrostatic not implemented for fv_subgridz"
grid_indexing = stencil_factory.grid_indexing
self._k_sponge = n_sponge
if self._k_sponge is not None:
if self._k_sponge < 3:
return
else:
self._k_sponge = grid_indexing.domain[2]
if self._k_sponge < min(grid_indexing.domain[2], 24):
t_max = T2_MAX
else:
t_max = T3_MAX
if nwat == 0:
xvir = 0.0
else:
xvir = ZVIR
self._m = 3
self._fv_sg_adj = float(fv_sg_adj)
self._is = grid_indexing.isc
self._js = grid_indexing.jsc
kbot_domain = (grid_indexing.domain[0], grid_indexing.domain[1], self._k_sponge)
origin = grid_indexing.origin_compute()
# TODO(eddied): Applying the `KCacheDetection` pass causes validation to fail
self._init_stencil = stencil_factory.from_origin_domain(
init,
origin=origin,
domain=(
grid_indexing.domain[0],
grid_indexing.domain[1],
self._k_sponge + 1,
),
skip_passes=("KCacheDetection",),
)
self._m_loop_stencil = stencil_factory.from_origin_domain(
m_loop,
externals={"t_max": t_max, "xvir": xvir},
origin=origin,
domain=kbot_domain,
)
self._finalize_stencil = stencil_factory.from_origin_domain(
finalize,
externals={
"hydrostatic": hydrostatic,
"fv_sg_adj": fv_sg_adj,
},
origin=origin,
domain=kbot_domain,
)
shape = grid_indexing.domain_full(add=(1, 1, 0))
self._q0 = {}
for tracername in utils.tracer_variables:
self._q0[tracername] = utils.make_storage_from_shape(shape)
self._tmp_u0 = utils.make_storage_from_shape(shape)
self._tmp_v0 = utils.make_storage_from_shape(shape)
self._tmp_w0 = utils.make_storage_from_shape(shape)
self._tmp_gz = utils.make_storage_from_shape(shape)
self._tmp_t0 = utils.make_storage_from_shape(shape)
self._tmp_static_energy = utils.make_storage_from_shape(shape)
self._tmp_total_energy = utils.make_storage_from_shape(shape)
self._tmp_cvm = utils.make_storage_from_shape(shape)
self._tmp_cpm = utils.make_storage_from_shape(shape)
self._ratios = {0: 0.25, 1: 0.5, 2: 0.999}
def __init__(
self,
stencil_factory: StencilFactory,
grid_data: GridData,
damping_coefficients: DampingCoefficients,
nested: bool,
stretched_grid: bool,
dddmp,
d4_bg,
nord,
grid_type,
nord_col: FloatFieldK,
d2_bg: FloatFieldK,
):
self.grid_indexing = stencil_factory.grid_indexing
assert not nested, "nested not implemented"
assert grid_type < 3, "Not implemented, grid_type>=3, specifically smag_corner"
# TODO: make dddmp a compile-time external, instead of runtime scalar
self._dddmp = dddmp
# TODO: make da_min_c a compile-time external, instead of runtime scalar
self._da_min_c = damping_coefficients.da_min_c
self._grid_type = grid_type
self._nord_column = nord_col
self._d2_bg_column = d2_bg
self._rarea_c = grid_data.rarea_c
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
self._cosa_u = grid_data.cosa_u
self._cosa_v = grid_data.cosa_v
self._sina_u = grid_data.sina_u
self._sina_v = grid_data.sina_v
self._dxc = grid_data.dxc
self._dyc = grid_data.dyc
# TODO: calculate these locally based on grid_data
self._divg_u = spec.grid.divg_u
self._divg_v = spec.grid.divg_v
nonzero_nord_k = 0
self._nonzero_nord = int(nord)
for k in range(len(self._nord_column)):
if self._nord_column[k] > 0:
nonzero_nord_k = k
self._nonzero_nord = int(self._nord_column[k])
break
if stretched_grid:
self._dd8 = damping_coefficients.da_min * d4_bg ** (self._nonzero_nord + 1)
else:
self._dd8 = (damping_coefficients.da_min_c * d4_bg) ** (
self._nonzero_nord + 1
)
kstart = nonzero_nord_k
nk = self.grid_indexing.domain[2] - kstart
self._do_zero_order = nonzero_nord_k > 0
low_k_stencil_factory = stencil_factory.restrict_vertical(
k_start=0, nk=nonzero_nord_k
)
high_k_stencil_factory = stencil_factory.restrict_vertical(
k_start=nonzero_nord_k
)
self.a2b_ord4 = AGrid2BGridFourthOrder(
stencil_factory=high_k_stencil_factory,
grid_data=grid_data,
grid_type=self._grid_type,
replace=False,
)
self._get_delpc = low_k_stencil_factory.from_dims_halo(
func=get_delpc,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(0, 0),
skip_passes=("GreedyMerging",),
)
self._damping = low_k_stencil_factory.from_dims_halo(
damping,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(0, 0),
)
self._copy_computeplus = high_k_stencil_factory.from_dims_halo(
func=basic.copy_defn,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(0, 0),
)
origins = []
origins_v = []
origins_u = []
domains = []
domains_v = []
domains_u = []
for n in range(1, self._nonzero_nord + 1):
nt = self._nonzero_nord - n
nint = self.grid_indexing.domain[0] + 2 * nt + 1
njnt = self.grid_indexing.domain[1] + 2 * nt + 1
js = self.grid_indexing.jsc - nt
is_ = self.grid_indexing.isc - nt
origins_v.append((is_ - 1, js, kstart))
domains_v.append((nint + 1, njnt, nk))
origins_u.append((is_, js - 1, kstart))
domains_u.append((nint, njnt + 1, nk))
origins.append((is_, js, kstart))
domains.append((nint, njnt, nk))
self._vc_from_divg_stencils = get_stencils_with_varied_bounds(
vc_from_divg,
origins=origins_v,
domains=domains_v,
stencil_factory=stencil_factory,
)
self._uc_from_divg_stencils = get_stencils_with_varied_bounds(
uc_from_divg,
origins=origins_u,
domains=domains_u,
stencil_factory=stencil_factory,
)
self._redo_divg_d_stencils = get_stencils_with_varied_bounds(
redo_divg_d,
origins=origins,
domains=domains,
stencil_factory=stencil_factory,
externals={"do_adjustment": not stretched_grid},
)
self._damping_nord_highorder_stencil = high_k_stencil_factory.from_dims_halo(
func=damping_nord_highorder_stencil,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(0, 0),
)
self._smagorinksy_diffusion_approx_stencil = (
high_k_stencil_factory.from_dims_halo(
func=smagorinksy_diffusion_approx,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(0, 0),
)
)
self._set_value = high_k_stencil_factory.from_dims_halo(
func=basic.set_value_defn,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(self.grid_indexing.n_halo, self.grid_indexing.n_halo),
)
self._corner_tmp = utils.make_storage_from_shape(self.grid_indexing.max_shape)
self.fill_corners_bgrid_x = corners.FillCornersBGrid(
direction="x",
temporary_field=self._corner_tmp,
stencil_factory=high_k_stencil_factory,
)
self.fill_corners_bgrid_y = corners.FillCornersBGrid(
direction="y",
temporary_field=self._corner_tmp,
stencil_factory=high_k_stencil_factory,
)
self._fill_corners_dgrid_stencil = high_k_stencil_factory.from_dims_halo(
func=corners.fill_corners_dgrid_defn,
compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
compute_halos=(self.grid_indexing.n_halo, self.grid_indexing.n_halo),
)
def __init__(
self,
stencil_factory: StencilFactory,
damping_coefficients: DampingCoefficients,
rarea,
nmax: int,
):
"""
Args:
grid: fv3core grid object
"""
grid_indexing = stencil_factory.grid_indexing
self._del6_u = damping_coefficients.del6_u
self._del6_v = damping_coefficients.del6_v
self._rarea = rarea
self._fx = utils.make_storage_from_shape(grid_indexing.max_shape)
self._fy = utils.make_storage_from_shape(grid_indexing.max_shape)
self._q = utils.make_storage_from_shape(grid_indexing.max_shape)
self._corner_fill = stencil_factory.from_dims_halo(
func=corner_fill,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._copy_stencil = stencil_factory.from_dims_halo(
func=copy_defn,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
compute_halos=(3, 3),
)
self._ntimes = min(3, nmax)
origins = []
domains_x = []
domains_y = []
domains = []
for n_halo in range(self._ntimes - 1, -1, -1):
origin, domain = grid_indexing.get_origin_domain(
[X_DIM, Y_DIM, Z_DIM], halos=(n_halo, n_halo))
_, domain_x = grid_indexing.get_origin_domain(
[X_INTERFACE_DIM, Y_DIM, Z_DIM], halos=(n_halo, n_halo))
_, domain_y = grid_indexing.get_origin_domain(
[X_DIM, Y_INTERFACE_DIM, Z_DIM], halos=(n_halo, n_halo))
origins.append(cast_to_index3d(origin))
domains.append(cast_to_index3d(domain))
domains_x.append(cast_to_index3d(domain_x))
domains_y.append(cast_to_index3d(domain_y))
self._compute_zonal_flux = get_stencils_with_varied_bounds(
compute_zonal_flux,
origins,
domains_x,
stencil_factory=stencil_factory)
self._compute_meridional_flux = get_stencils_with_varied_bounds(
compute_meridional_flux,
origins,
domains_y,
stencil_factory=stencil_factory)
self._update_q = get_stencils_with_varied_bounds(
update_q, origins, domains, stencil_factory=stencil_factory)
self._copy_corners_x: corners.CopyCorners = corners.CopyCorners(
direction="x", stencil_factory=stencil_factory)
"""Stencil responsible for doing corners updates in x-direction."""
self._copy_corners_y: corners.CopyCorners = corners.CopyCorners(
direction="y", stencil_factory=stencil_factory)
"""Stencil responsible for doing corners updates in y-direction."""
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,
stencil_factory: StencilFactory,
grid_data: GridData,
damping_coefficients: DampingCoefficients,
grid_type: int,
hord,
nord=None,
damp_c=None,
):
# use a shorter alias for grid_indexing here to avoid very verbose lines
idx = stencil_factory.grid_indexing
self._area = grid_data.area
origin = idx.origin_compute()
self._q_advected_y = utils.make_storage_from_shape(
idx.max_shape, origin)
self._q_advected_x = utils.make_storage_from_shape(
idx.max_shape, origin)
self._q_x_advected_mean = utils.make_storage_from_shape(
idx.max_shape, origin)
self._q_y_advected_mean = utils.make_storage_from_shape(
idx.max_shape, origin)
self._q_advected_x_y_advected_mean = utils.make_storage_from_shape(
idx.max_shape, origin)
self._q_advected_y_x_advected_mean = utils.make_storage_from_shape(
idx.max_shape, origin)
self._corner_tmp = utils.make_storage_from_shape(
idx.max_shape, origin=idx.origin_full())
"""Temporary field to use for corner computation in both x and y direction"""
self._nord = nord
self._damp_c = damp_c
ord_outer = hord
ord_inner = 8 if hord == 10 else hord
self.q_i_stencil = stencil_factory.from_origin_domain(
q_i_stencil,
origin=idx.origin_full(add=(0, 3, 0)),
domain=idx.domain_full(add=(0, -3, 1)),
)
self.q_j_stencil = stencil_factory.from_origin_domain(
q_j_stencil,
origin=idx.origin_full(add=(3, 0, 0)),
domain=idx.domain_full(add=(-3, 0, 1)),
)
self.stencil_transport_flux = stencil_factory.from_origin_domain(
final_fluxes,
origin=idx.origin_compute(),
domain=idx.domain_compute(add=(1, 1, 1)),
)
if (self._nord is not None) and (self._damp_c is not None):
self.delnflux: Optional[DelnFlux] = DelnFlux(
stencil_factory=stencil_factory,
damping_coefficients=damping_coefficients,
rarea=grid_data.rarea,
nord=self._nord,
damp_c=self._damp_c,
)
else:
self.delnflux = None
self.x_piecewise_parabolic_inner = XPiecewiseParabolic(
stencil_factory=stencil_factory,
dxa=grid_data.dxa,
grid_type=grid_type,
iord=ord_inner,
origin=idx.origin_compute(add=(0, -idx.n_halo, 0)),
domain=idx.domain_compute(add=(1, 1 + 2 * idx.n_halo, 1)),
)
self.y_piecewise_parabolic_inner = YPiecewiseParabolic(
stencil_factory=stencil_factory,
dya=grid_data.dya,
grid_type=grid_type,
jord=ord_inner,
origin=idx.origin_compute(add=(-idx.n_halo, 0, 0)),
domain=idx.domain_compute(add=(1 + 2 * idx.n_halo, 1, 1)),
)
self.x_piecewise_parabolic_outer = XPiecewiseParabolic(
stencil_factory=stencil_factory,
dxa=grid_data.dxa,
grid_type=grid_type,
iord=ord_outer,
origin=idx.origin_compute(),
domain=idx.domain_compute(add=(1, 1, 1)),
)
self.y_piecewise_parabolic_outer = YPiecewiseParabolic(
stencil_factory=stencil_factory,
dya=grid_data.dya,
grid_type=grid_type,
jord=ord_outer,
origin=idx.origin_compute(),
domain=idx.domain_compute(add=(1, 1, 1)),
)
def __init__(
self,
stencil_factory: StencilFactory,
kord: int,
iv: int,
i1: int,
i2: int,
j1: int,
j2: int,
):
"""
The constraints on the spline are set by kord and iv.
Arguments:
stencil_factory
kord: ???
iv: ???
i1: The first i-element to compute on
i2: The last i-element to compute on
j1: The first j-element to compute on
j2: The last j-element to compute on
"""
assert kord <= 10, f"kord {kord} not implemented."
grid_indexing = stencil_factory.grid_indexing
full_orig: Tuple[int] = grid_indexing.origin_full()
km: int = grid_indexing.domain[2]
self._kord = kord
self._gam: FloatField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig
)
self._q: FloatField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig
)
self._q_bot: FloatField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig
)
self._extm: BoolField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig, dtype=bool
)
self._ext5: BoolField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig, dtype=bool
)
self._ext6: BoolField = utils.make_storage_from_shape(
grid_indexing.domain_full(add=(0, 0, 1)), origin=full_orig, dtype=bool
)
i_extent: int = i2 - i1 + 1
j_extent: int = j2 - j1 + 1
origin: Tuple[int, int, int] = (i1, j1, 0)
domain: Tuple[int, int, int] = (i_extent, j_extent, km)
domain_extend: Tuple[int, int, int] = (i_extent, j_extent, km + 1)
self._set_initial_values = stencil_factory.from_origin_domain(
func=set_initial_vals,
externals={"iv": iv, "kord": abs(kord)},
origin=origin,
domain=domain_extend,
)
self._apply_constraints = stencil_factory.from_origin_domain(
func=apply_constraints,
externals={"iv": iv, "kord": abs(kord)},
origin=origin,
domain=domain,
)
self._set_interpolation_coefficients = stencil_factory.from_origin_domain(
func=set_interpolation_coefficients,
externals={"iv": iv, "kord": abs(kord)},
origin=origin,
domain=domain,
)