def test_stencil_without_effect(backend):
def definition1(field_in: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
tmp = 0.0
def definition2(f_in: gtscript.Field[np.float_]):
from __externals__ import flag
with computation(PARALLEL), interval(...):
if __INLINED(flag):
B = f_in
stencil1 = gtscript.stencil(backend, definition1)
stencil2 = gtscript.stencil(backend, definition2, externals={"flag": False})
field_in = gt_storage.ones(
dtype=np.float_, backend=backend, shape=(23, 23, 23), default_origin=(0, 0, 0)
)
# test with explicit domain specified
stencil1(field_in, domain=(3, 3, 3))
stencil2(field_in, domain=(3, 3, 3))
# test without domain specified
stencil1(field_in)
def test_negative_origin(backend):
def stencil_i(
input_field: gtscript.Field[gtscript.IJK, np.int32],
output_field: gtscript.Field[gtscript.IJK, np.int32],
):
with computation(PARALLEL), interval(...):
output_field = input_field[1, 0, 0]
def stencil_k(
input_field: gtscript.Field[gtscript.IJK, np.int32],
output_field: gtscript.Field[gtscript.IJK, np.int32],
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 1]
input_field = gt_storage.ones(
backend, default_origin=(0, 0, 0), shape=(1, 1, 1), dtype=np.int32
)
output_field = gt_storage.zeros(
backend, default_origin=(0, 0, 0), shape=(1, 1, 1), dtype=np.int32
)
for origin, stencil in (((-1, 0, 0), stencil_i), ((0, 0, -1), stencil_k)):
gtscript.stencil(definition=stencil, backend=backend)(
input_field, output_field, origin={"input_field": origin}
)
assert output_field[0, 0, 0] == 1
def test_default_arguments(backend):
branch_true = gtscript.stencil(backend=backend,
definition=a_stencil,
externals={"BRANCH": True},
rebuild=True)
branch_false = gtscript.stencil(backend=backend,
definition=a_stencil,
externals={"BRANCH": False},
rebuild=True)
arg1 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg2 = gt_storage.zeros(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg3 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
tmp = np.asarray(arg3)
tmp *= 2
branch_true(arg1, None, arg3, par1=2.0)
np.testing.assert_equal(arg1, 14 * np.ones((3, 3, 3)))
branch_true(arg1, None, par1=2.0)
np.testing.assert_equal(arg1, 196 * np.ones((3, 3, 3)))
branch_false(arg1, arg2, arg3, par1=2.0, par3=2.0)
np.testing.assert_equal(arg1, 56 * np.ones((3, 3, 3)))
with pytest.raises((ValueError, AssertionError)):
branch_false(arg1, arg2, par1=2.0, par3=2.0)
arg1 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg2 = gt_storage.zeros(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg3 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
tmp = np.asarray(arg3)
tmp *= 2
branch_true(arg1, arg2=None, par1=2.0, par2=5.0, par3=3.0)
np.testing.assert_equal(arg1, 10 * np.ones((3, 3, 3)))
branch_true(arg1, arg2=None, par1=2.0, par2=5.0)
np.testing.assert_equal(arg1, 100 * np.ones((3, 3, 3)))
branch_false(arg1, arg2, arg3, par1=2.0, par2=5.0, par3=3.0)
np.testing.assert_equal(arg1, 60 * np.ones((3, 3, 3)))
with pytest.raises((TypeError, AssertionError)):
branch_false(arg1, arg2, arg3, par1=2.0, par2=5.0)
def test_exec_info(backend):
"""test that proper warnings are raised depending on field type."""
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
exec_info = {}
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(1, 1, 0),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend, shape=(23, 23, 10), default_origin=(1, 1, 0), dtype=np.float64
)
stencil(
in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10),
exec_info=exec_info,
)
timings = ["call", "call_run", "run"]
assert all([k + "_start_time" in exec_info for k in timings])
assert all([k + "_end_time" in exec_info for k in timings])
assert all([exec_info[k + "_end_time"] > exec_info[k + "_start_time"] for k in timings])
if backend.startswith("gt:"):
assert "run_cpp_start_time" in exec_info
assert "run_cpp_end_time" in exec_info
assert exec_info["run_cpp_end_time"] > exec_info["run_cpp_start_time"]
def test_ndarray_warning():
"""test that proper warnings are raised depending on field type."""
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
with pytest.warns(RuntimeWarning):
stencil(
in_field=in_field.view(np.ndarray),
out_field=out_field.view(np.ndarray),
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
with pytest.warns(None) as record:
stencil(
in_field=in_field,
out_field=out_field,
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
assert len(record) == 0
def test_generation(self, test, externals_dict):
"""Test source code generation for all *backends* and *stencil suites*.
The generated implementations are cached in a :class:`utils.ImplementationsDB`
instance, to avoid duplication of (potentially expensive) compilations.
"""
cls = type(self)
implementation = gtscript.stencil(
backend=test["backend"],
definition=test["definition"],
name=f"{test['suite']}_{test['backend']}_{test['test_id']}",
rebuild=True,
externals=externals_dict,
)
for k, v in externals_dict.items():
implementation._gt_constants_[k] = v
assert isinstance(implementation, StencilObject)
assert implementation.backend == test["backend"]
assert all(
cls.global_boundaries[name] == field_info.boundary
for name, field_info in implementation._gt_field_info_.items()
)
test["implementations"].append(implementation)
def test_generation(self, test, externals_dict):
"""Test source code generation for all *backends* and *stencil suites*.
The generated implementations are cached in a :class:`utils.ImplementationsDB`
instance, to avoid duplication of (potentially expensive) compilations.
"""
cls = type(self)
backend_slug = gt_utils.slugify(test["backend"], valid_symbols="")
implementation = gtscript.stencil(
backend=test["backend"],
definition=test["definition"],
name=f"{test['suite']}_{backend_slug}_{test['test_id']}",
rebuild=True,
externals=externals_dict,
# debug_mode=True,
# _impl_opts={"cache-validation": False, "code-generation": False},
)
for k, v in externals_dict.items():
implementation.constants[k] = v
assert isinstance(implementation, StencilObject)
assert implementation.backend == test["backend"]
assert all(field_info.boundary >= cls.global_boundaries[name]
for name, field_info in implementation.field_info.items())
test["implementations"].append(implementation)
def test_compilation(self, dtype_in, dtype_out, dtype_scalar):
definition = self.sumdiff_defs
dtypes = {
"dtype_in": dtype_in,
"dtype_out": dtype_out,
"dtype_scalar": dtype_scalar
}
sumdiff = gtscript.stencil("debug", definition, dtypes=dtypes)
annotations = getattr(definition, "__annotations__", {})
assert "in_a" in annotations
assert isinstance(annotations["in_a"], gtscript._FieldDescriptor)
assert annotations["in_a"].dtype == "dtype_in"
assert "in_b" in annotations
assert isinstance(annotations["in_b"], gtscript._FieldDescriptor)
assert annotations["in_b"].dtype == "dtype_in"
assert "out_c" in annotations
assert isinstance(annotations["out_c"], gtscript._FieldDescriptor)
assert annotations["out_c"].dtype == "dtype_out"
assert "out_d" in annotations
assert isinstance(annotations["out_d"], gtscript._FieldDescriptor)
assert annotations["out_d"].dtype == float
assert "wa" in annotations
assert annotations["wa"] == "dtype_scalar"
assert "wb" in annotations
assert annotations["wb"] == int
assert len(annotations) == 6
def _test_generation(cls, test, externals_dict):
"""Test source code generation for all *backends* and *stencil suites*.
The generated implementations are cached in a :class:`utils.ImplementationsDB`
instance, to avoid duplication of (potentially expensive) compilations.
"""
backend_slug = gt_utils.slugify(test["backend"], valid_symbols="")
implementation = gtscript.stencil(
backend=test["backend"],
definition=test["definition"],
name=cls.__module__ +
f".{test['suite']}_{backend_slug}_{test['test_id']}",
rebuild=True,
externals=externals_dict,
)
for k, v in externals_dict.items():
implementation.constants[k] = v
assert isinstance(implementation, StencilObject)
assert implementation.backend == test["backend"]
# Assert strict equality for Dawn backends
if implementation.backend.startswith("dawn"):
assert all(
field_info.boundary == cls.global_boundaries[name]
for name, field_info in implementation.field_info.items()
if field_info is not None)
else:
assert all(
field_info.boundary >= cls.global_boundaries[name]
for name, field_info in implementation.field_info.items()
if field_info is not None)
test["implementations"].append(implementation)
def _test_generation(cls, test, externals_dict):
"""Test source code generation for all *backends* and *stencil suites*.
The generated implementations are cached in a :class:`utils.ImplementationsDB`
instance, to avoid duplication of (potentially expensive) compilations.
"""
backend_slug = gt_utils.slugify(test["backend"], valid_symbols="")
implementation = gtscript.stencil(
backend=test["backend"],
definition=test["definition"],
name=cls.__module__ +
f".{test['suite']}_{backend_slug}_{test['test_id']}",
rebuild=True,
externals=externals_dict,
)
for k, v in externals_dict.items():
implementation.constants[k] = v
assert isinstance(implementation, StencilObject)
assert implementation.backend == test["backend"]
for name, field_info in implementation.field_info.items():
if field_info.access == AccessKind.NONE:
continue
for i, ax in enumerate("IJK"):
assert (ax not in field_info.axes or ax == "K"
or field_info.boundary[i] >=
cls.global_boundaries[name][i])
test["implementations"].append(implementation)
def test_halo_checks(backend):
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test default works
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field, out_field=out_field)
assert (out_field[1:-1, 1:-1, :] == 1).all()
# test setting arbitrary, small domain works
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(10, 10, 10))
assert (out_field[2:12, 2:12, :] == 1).all()
# test setting domain+origin too large raises
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
with pytest.raises(ValueError):
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10))
# test 2*origin+domain does not raise if still fits (c.f. previous bug in c++ check.)
in_field = gt_storage.ones(backend=backend,
shape=(23, 23, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(23, 23, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10))
def test_assert_same_shape(backend):
stencil_call = gtscript.stencil(definition=stencil, backend=backend)
A = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
B = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(2, 2, 2))
C = gt_storage.ones(backend=backend,
dtype=np.float32,
shape=(3, 3, 3),
default_origin=(0, 1, 0))
stencil_call(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
A = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(5, 5, 5),
default_origin=(0, 0, 0))
A = A[1:-1, 1:-1, 1:-1]
A.is_stencil_view = True
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
C = gt_storage.ones(backend=backend,
dtype=np.float32,
shape=(5, 5, 5),
default_origin=(0, 1, 0))
with pytest.raises(ValueError):
stencil_call(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
with pytest.raises(ValueError):
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
def test_read_data_dim_indirect_addressing(backend):
INT32_VEC2 = (np.int32, (2,))
def stencil(
input_field: gtscript.Field[gtscript.IJK, INT32_VEC2],
output_field: gtscript.Field[gtscript.IJK, np.int32],
index: int,
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 0][index]
default_origin = (0, 0, 0)
full_shape = (1, 1, 2)
input_field = gt_storage.ones(backend, default_origin, full_shape, dtype=INT32_VEC2)
output_field = gt_storage.zeros(backend, default_origin, full_shape, dtype=np.int32)
gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, 1)
assert output_field[0, 0, 0] == 1
def test_origin_selection():
stencil = gtscript.stencil(definition=base_stencil, backend="numpy")
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
assert A[1, 1, 1] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin={"_all_": (1, 1, 1), "field1": (2, 2, 2)}, domain=(1, 1, 1))
assert A[2, 2, 2] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin={"field1": (2, 2, 2)}, domain=(1, 1, 1))
assert A[2, 2, 2] == 4
assert B[2, 2, 2] == 7
assert C[0, 1, 0] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
def run_stencil_scaling(stencil_def,
backends,
domain,
origin,
nruns=5,
factor=2,
dtype=np.float64):
ni, nj, nk = domain
nhalo = origin[0]
init_ni = ni
init_nj = nj
timings = dict()
sizes = dict()
for backend in backends:
ni = init_ni
nj = init_nj
sizes[backend] = []
timings[backend] = []
for n in range(0, nruns):
print(f"Running with {backend} backend ({n})...")
domain = (ni, nj, nk)
shape = (ni + 2 * nhalo, nj + 2 * nhalo, nk)
rand_data = np.random.randn(*shape)
in_field = gt4py.storage.from_array(rand_data,
backend,
origin,
shape,
dtype=dtype)
out_field = gt4py.storage.zeros(backend, origin, shape, dtype)
exec_info = {}
stencil = gtscript.stencil(backend, stencil_def)
stencil(in_field,
out_field,
origin=origin,
domain=domain,
exec_info=exec_info)
call_time = (exec_info['call_end_time'] -
exec_info['call_start_time']) * 1000.
run_time = (exec_info['run_end_time'] -
exec_info['run_start_time']) * 1000.
timings[backend].append(run_time)
sizes[backend].append(ni)
ni *= factor
nj *= factor
return timings, sizes
def corner_ut(
uc,
vc,
ut,
vt,
cosa_u,
cosa_v,
ui,
uj,
vi,
vj,
west,
lower,
south=True,
vswitch=False,
):
if vswitch:
lowerfactor = 1 if lower else -1
else:
lowerfactor = 1
vx = vi + index_offset(west, False, south) * lowerfactor
ux = ui + index_offset(west, True, south) * lowerfactor
vy = vj + index_offset(lower, False, south) * lowerfactor
uy = uj + index_offset(lower, True, south) * lowerfactor
if stencil_corner:
decorator = gtscript.stencil(
backend=global_config.get_backend(),
externals={
"vi": vi - ui,
"vj": vj - uj,
"ux": ux - ui,
"uy": uy - uj,
"vx": vx - ui,
"vy": vy - uj,
},
rebuild=global_config.get_rebuild(),
)
corner_stencil = decorator(corner_ut_stencil)
corner_stencil(
uc,
vc,
ut,
vt,
cosa_u,
cosa_v,
origin=(ui, uj, 0),
domain=(1, 1, grid().npz),
)
else:
damp = get_damp(cosa_u, cosa_v, ui, uj, vi, vj)
ut[ui,
uj, :] = (uc[ui, uj, :] - 0.25 * cosa_u[ui, uj, :] *
(vt[vi, vy, :] + vt[vx, vy, :] + vt[vx, vj, :] +
vc[vi, vj, :] - 0.25 * cosa_v[vi, vj, :] *
(ut[ux, uj, :] + ut[ux, uy, :] + ut[ui, uy, :]))) * damp
def run(in_dict, backend):
"""Run function for GFS thermodynamics surface ice model
With this function, the GFS thermodynamics surface ice model can be run
as a standalone parameterization.
"""
# TODO - remove this one we have high-dimensional fields
# special handling of stc
stc = in_dict.pop("stc")
in_dict["stc0"] = stc[:, 0]
in_dict["stc1"] = stc[:, 1]
# setup storages
scalar_dict = {k: in_dict[k] for k in SCALAR_VARS}
out_dict = {
k: numpy_to_gt4py_storage(in_dict[k].copy(), backend=backend)
for k in OUT_VARS
}
in_dict = {
k: numpy_to_gt4py_storage(in_dict[k], backend=backend)
for k in IN_VARS
}
# compile stencil
sfc_sice = gtscript.stencil(definition=sfc_sice_defs,
backend=backend,
externals={})
# set timer
tic = timeit.default_timer()
# call sea-ice parametrization
sfc_sice(**in_dict, **out_dict, **scalar_dict)
# set timer
toc = timeit.default_timer()
# calculate elapsed time
elapsed_time = toc - tic
# convert back to numpy for validation
out_dict = {
k: gt4py_to_numpy_storage(out_dict[k], backend=backend)
for k in OUT_VARS
}
# TODO - remove this one we have high-dimensional fields
# special handling of stc
stc[:, 0] = out_dict.pop("stc0")[:]
stc[:, 1] = out_dict.pop("stc1")[:]
out_dict["stc"] = stc
return out_dict, elapsed_time
def build(self, **kwargs):
"""Build the stencil, get ready to execute."""
externals = self.externals()
externals.update(kwargs.pop("externals", {}))
self._stencil = gtscript.stencil(
definition=self.stencil_definition(),
backend=self._backend,
externals=externals,
**kwargs,
)
return self
def test_read_data_dim_indirect_addressing(backend):
INT32_VEC2 = (np.int32, (2,))
def stencil(
input_field: gtscript.Field[gtscript.IJK, INT32_VEC2],
output_field: gtscript.Field[gtscript.IJK, np.int32],
index: int,
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 0][index]
default_origin = (0, 0, 0)
full_shape = (1, 1, 2)
input_field = gt_storage.ones(backend, default_origin, full_shape, dtype=INT32_VEC2)
output_field = gt_storage.zeros(backend, default_origin, full_shape, dtype=np.int32)
if backend in (backend.values[0] for backend in LEGACY_GRIDTOOLS_BACKENDS):
with pytest.raises(ValueError):
gtscript.stencil(definition=stencil, backend=backend)
else:
gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, 1)
assert output_field[0, 0, 0] == 1
def do_test(data_file, backend):
data = Dataset(data_file, backend)
# other fields
pe = data.new(IJK, float, pad_k=True)
riem = stencil(backend=backend,
definition=riem_solver_c,
externals={"A_IMP": data["a_imp"]})
riem(data["ms"], data["dt"], data["akap"], data["cappa"], data["cp"],
data["ptop"], data["hs"], data["w3"], data["pt"], data["q_con"],
data["delp"], data["gz"], data["pef"], data["ws"], data["p_fac"],
data["scale_m"], pe)
def test_k_bounds_exec(definition, expected):
expected_k_bounds, expected_min_k_size = expected["k_bounds"], expected[
"min_k_size"]
required_field_size = expected_min_k_size + expected_k_bounds[
0] + expected_k_bounds[1]
if required_field_size > 0:
backend = "gtc:gt:cpu_ifirst"
compiled_stencil = stencil(backend, definition)
field_a = gt4py.storage.zeros(
backend=backend,
default_origin=(0, 0, 0),
shape=(1, 1, expected_min_k_size),
dtype=np.float64,
)
field_b = gt4py.storage.ones(
backend=backend,
default_origin=(0, 0, 0),
shape=(1, 1, required_field_size),
dtype=np.float64,
)
# test with correct domain, origin to low
with pytest.raises(ValueError,
match="Origin for field field_b too small"):
compiled_stencil(
field_a,
field_b,
domain=(1, 1, expected_min_k_size),
origin={"field_b": (0, 0, expected_k_bounds[0] - 1)},
)
# test with correct domain, correct origin
compiled_stencil(
field_a,
field_b,
domain=(1, 1, expected_min_k_size),
origin={"field_b": (0, 0, expected_k_bounds[0])},
)
# test with wrong domain, correct origin
with pytest.raises(ValueError,
match="Compute domain too small. Sequential axis"):
compiled_stencil(
field_a,
field_b,
domain=(1, 1, expected_min_k_size - 1),
origin={"field_b": (0, 0, expected_k_bounds[0])},
)
def generate_test_module(name, backend, *, id_version, rebuild=True):
module_name = "_test_module." + name
stencil_name = name
backend_opts = {}
if issubclass(backend, gt_back.BaseGTBackend):
backend_opts["debug_mode"] = False
backend_opts["add_profile_info"] = True
backend_opts["verbose"] = True
options = gt_defs.BuildOptions(
name=stencil_name, module=module_name, rebuild=rebuild, backend_opts=backend_opts
)
decorator = gtscript.stencil(backend=backend.name, externals=EXTERNALS_REGISTRY[stencil_name])
stencil_definition = stencil_registry[name]
return decorator(stencil_definition)
def test_generation(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend, stencil_definition, externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=(v.dtype, v.data_dims) if v.data_dims else v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
# vertical domain size >= 16 required for test_large_k_interval
stencil(**args, origin=(10, 10, 5), domain=(3, 3, 16))
def test_generation_gpu(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend,
stencil_definition,
externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
stencil(**args, origin=(10, 10, 10), domain=(3, 3, 3))
def test_domain_selection():
stencil = gtscript.stencil(definition=base_stencil, backend="numpy")
A = gt_storage.ones(backend="gtmc",
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
B = gt_storage.ones(backend="gtx86",
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(2, 2, 2))
C = gt_storage.ones(backend="numpy",
dtype=np.float32,
shape=(3, 3, 3),
default_origin=(0, 1, 0))
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
assert A[1, 1, 1] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(backend="gtmc",
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
B = gt_storage.ones(backend="gtx86",
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(2, 2, 2))
C = gt_storage.ones(backend="numpy",
dtype=np.float32,
shape=(3, 3, 3),
default_origin=(0, 1, 0))
stencil(A, B, C, param=3.0, origin=(0, 0, 0))
assert np.all(A == 4)
assert np.all(B == 7)
assert np.all(C == 21)
def test_np_int_types():
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(np.int64(1), int(1), 0),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(np.int64(1), int(1), 0),
dtype=np.float64,
)
stencil(
in_field=in_field,
out_field=out_field,
origin=(np.int8(2), np.int16(2), np.int32(0)),
domain=(np.int64(20), int(20), 10),
)
def test_np_array_int_types():
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
stencil(
in_field=in_field,
out_field=out_field,
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
def __call__(self, *args, origin: Index3D, domain: Index3D, **kwargs) -> None:
"""Call the stencil, compiling the stencil if necessary.
The stencil needs to be recompiled if any of the following changes
1. the origin and/or domain
2. any external value
3. the function signature or code
Args:
domain: Stencil compute domain (required)
origin: Data index mapped to (0, 0, 0) in the compute domain (required)
"""
# Can optimize this by marking stencils that need these
axis_offsets = fv3core.utils.axis_offsets(spec.grid, origin, domain)
regenerate_stencil = not self.built or global_config.get_rebuild()
if not regenerate_stencil:
# Check if we really do need to regenerate
for key, value in self.axis_offsets.items():
if axis_offsets[key] != value:
axis_offsets_changed = True
break
else:
axis_offsets_changed = False
regenerate_stencil = regenerate_stencil or axis_offsets_changed
if regenerate_stencil:
new_build_info = {}
stencil_kwargs = {
"rebuild": global_config.get_rebuild(),
"backend": global_config.get_backend(),
"externals": {
"namelist": spec.namelist,
"grid": spec.grid,
**axis_offsets,
**self._passed_externals,
},
"format_source": global_config.get_format_source(),
"build_info": new_build_info,
**self.backend_kwargs,
}
# gtscript.stencil always returns a new class instance even if it
# used the cached module.
self.stencil_object = gtscript.stencil(
definition=self.func, **stencil_kwargs
)
if self.stencil_object not in self._axis_offsets_cache:
def_ir = stencil_kwargs["build_info"]["def_ir"]
axis_offsets = {
k: v for k, v in def_ir.externals.items() if k in axis_offsets
}
self._axis_offsets_cache[self.stencil_object] = axis_offsets
# Call it
exec_info = {}
kwargs["exec_info"] = kwargs.get("exec_info", exec_info)
kwargs["validate_args"] = kwargs.get("validate_args", utils.validate_args)
name = f"{self.func.__module__}.{self.func.__name__}"
_maybe_save_report(
f"{name}-before",
self.times_called,
self.func.__dict__["_gtscript_"]["api_signature"],
args,
kwargs,
)
self.stencil_object(*args, **kwargs, origin=origin, domain=domain)
_maybe_save_report(
f"{name}-after",
self.times_called,
self.func.__dict__["_gtscript_"]["api_signature"],
args,
kwargs,
)
self.times_called += 1
# Update timers
exec_info = kwargs["exec_info"]
self.timers.run += exec_info["run_end_time"] - exec_info["run_start_time"]
self.timers.call_run += (
exec_info["call_run_end_time"] - exec_info["call_run_start_time"]
)
def compile_stencils(self, backend):
self.advection = gtscript.stencil(backend=backend, definition=self.advection_def)
self.diffusion = gtscript.stencil(backend=backend, definition=self.diffusion_def)
def test_daa_warn(definition):
backend = "gtc:gt:cpu_ifirst"
with pytest.warns(UserWarning, match="`tmp` may be uninitialized."):
stencil(backend, definition)
