class TestHorizontalDiffusionSubroutines3(gt_testing.StencilTestSuite): """Diffusion in a horizontal 2D plane . """ dtypes = (np.float_, ) domain_range = [(1, 15), (1, 15), (1, 15)] backends = ["debug", "numpy", "gtx86"] symbols = dict( fwd_diff=gt_testing.global_name(singleton=fwd_diff_op_xy_varargin), BRANCH=gt_testing.global_name(one_of=(False, )), u=gt_testing.field(in_range=(-10, 10), boundary=[(2, 2), (2, 2), (0, 0)]), diffusion=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), weight=gt_testing.parameter(in_range=(0, 0.5)), ) def definition(u, diffusion, *, weight): """ Horizontal diffusion stencil. Parameters ---------- u : 3D float field, input diffusion : 3D float field, output weight : diffusion coefficient """ from __externals__ import fwd_diff, BRANCH with computation(PARALLEL), interval(...): laplacian = lap_op(u=u) laplacian2 = lap_op(u=u) if __INLINED(BRANCH): flux_i, flux_j = fwd_diff(field1=laplacian, field2=laplacian2) else: flux_i = fwd_diff(field1=laplacian) flux_j = fwd_diff_op_y(field=laplacian) diffusion = u[0, 0, 0] - weight * (flux_i[0, 0, 0] - flux_i[-1, 0, 0] + flux_j[0, 0, 0] - flux_j[0, -1, 0]) def validation(u, diffusion, *, weight, domain, origin, **kwargs): laplacian = 4.0 * u[1:-1, 1:-1, :] - (u[2:, 1:-1, :] + u[:-2, 1:-1, :] + u[1:-1, 2:, :] + u[1:-1, :-2, :]) flux_i = laplacian[1:, 1:-1, :] - laplacian[:-1, 1:-1, :] flux_j = laplacian[1:-1, 1:, :] - laplacian[1:-1, :-1, :] diffusion[...] = u[ 2:-2, 2:-2, :] - weight * (flux_i[1:, :, :] - flux_i[:-1, :, :] + flux_j[:, 1:, :] - flux_j[:, :-1, :])
class TestTwoOptionalFields(gt_testing.StencilTestSuite): dtypes = (np.float_, ) domain_range = [(1, 32), (1, 32), (1, 32)] backends = INTERNAL_BACKENDS symbols = dict( PHYS_TEND_A=gt_testing.global_name(one_of=(False, True)), PHYS_TEND_B=gt_testing.global_name(one_of=(False, True)), in_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), in_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), out_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), out_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), dyn_tend_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), dyn_tend_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), phys_tend_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), phys_tend_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), dt=gt_testing.parameter(in_range=(0, 100)), ) definition = two_optional_fields def validation( in_a, in_b, out_a, out_b, dyn_tend_a, dyn_tend_b, phys_tend_a=None, phys_tend_b=None, *, dt, domain, origin, **kwargs, ): out_a[...] = in_a + dt * dyn_tend_a out_b[...] = in_b + dt * dyn_tend_b if PHYS_TEND_A: # noqa: F821 # Undefined name out_a += dt * phys_tend_a if PHYS_TEND_B: # noqa: F821 # Undefined name out_b += dt * phys_tend_b
class TestHorizontalDiffusionSubroutines(gt_testing.StencilTestSuite): """Diffusion in a horizontal 2D plane . """ dtypes = (np.float_,) domain_range = [(1, 15), (1, 15), (1, 15)] backends = CPU_BACKENDS symbols = dict( fwd_diff=gt_testing.global_name(singleton=wrap1arg2return), u=gt_testing.field(in_range=(-10, 10), boundary=[(2, 2), (2, 2), (0, 0)]), diffusion=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), weight=gt_testing.parameter(in_range=(0, 0.5)), ) def definition(u, diffusion, *, weight): from __externals__ import fwd_diff with computation(PARALLEL), interval(...): laplacian = lap_op(u=u) flux_i, flux_j = fwd_diff(field=laplacian) diffusion = u[0, 0, 0] - weight * ( flux_i[0, 0, 0] - flux_i[-1, 0, 0] + flux_j[0, 0, 0] - flux_j[0, -1, 0] ) def validation(u, diffusion, *, weight, domain, origin, **kwargs): laplacian = 4.0 * u[1:-1, 1:-1, :] - ( u[2:, 1:-1, :] + u[:-2, 1:-1, :] + u[1:-1, 2:, :] + u[1:-1, :-2, :] ) flux_i = laplacian[1:, 1:-1, :] - laplacian[:-1, 1:-1, :] flux_j = laplacian[1:-1, 1:, :] - laplacian[1:-1, :-1, :] diffusion[...] = u[2:-2, 2:-2, :] - weight * ( flux_i[1:, :, :] - flux_i[:-1, :, :] + flux_j[:, 1:, :] - flux_j[:, :-1, :] )
class TestOptionalField(gt_testing.StencilTestSuite): dtypes = (np.float_, ) domain_range = [(1, 32), (1, 32), (1, 32)] backends = INTERNAL_BACKENDS symbols = dict( PHYS_TEND=gt_testing.global_name(one_of=(False, True)), in_field=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), out_field=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), dyn_tend=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), phys_tend=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), dt=gt_testing.parameter(in_range=(0, 100)), ) definition = optional_field def validation(in_field, out_field, dyn_tend, phys_tend=None, *, dt, domain, origin, **kwargs): out_field[...] = in_field + dt * dyn_tend if PHYS_TEND: # noqa: F821 # Undefined name out_field += dt * phys_tend
class TestHorizontalDiffusionSubroutines2(gt_testing.StencilTestSuite): """Diffusion in a horizontal 2D plane .""" dtypes = (np.float_, ) domain_range = [(1, 15), (1, 15), (1, 15)] backends = INTERNAL_BACKENDS symbols = dict( fwd_diff=gt_testing.global_name(singleton=fwd_diff_op_xy), BRANCH=gt_testing.global_name(one_of=(True, False)), u=gt_testing.field(in_range=(-10, 10), boundary=[(2, 2), (2, 2), (0, 0)]), diffusion=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), weight=gt_testing.parameter(in_range=(0, 0.5)), ) def definition(u, diffusion, *, weight): from __externals__ import BRANCH from __gtscript__ import __INLINED with computation(PARALLEL), interval(...): laplacian = lap_op(u=u) if __INLINED(BRANCH): flux_i = fwd_diff_op_x(field=laplacian) flux_j = fwd_diff_op_y(field=laplacian) else: flux_i, flux_j = fwd_diff_op_xy(field=laplacian) diffusion = u[ # noqa: F841 # Local name is assigned to but never used 0, 0, 0] - weight * (flux_i[0, 0, 0] - flux_i[-1, 0, 0] + flux_j[0, 0, 0] - flux_j[0, -1, 0]) def validation(u, diffusion, *, weight, domain, origin, **kwargs): laplacian = 4.0 * u[1:-1, 1:-1, :] - (u[2:, 1:-1, :] + u[:-2, 1:-1, :] + u[1:-1, 2:, :] + u[1:-1, :-2, :]) flux_i = laplacian[1:, 1:-1, :] - laplacian[:-1, 1:-1, :] flux_j = laplacian[1:-1, 1:, :] - laplacian[1:-1, :-1, :] diffusion[...] = u[ 2:-2, 2:-2, :] - weight * (flux_i[1:, :, :] - flux_i[:-1, :, :] + flux_j[:, 1:, :] - flux_j[:, :-1, :])
class TestParametricMix(gt_testing.StencilTestSuite): """Linear combination of input fields using several parameters. """ dtypes = { ("USE_ALPHA", ): np.int_, ("field_a", "field_b", "field_c"): np.float64, ("field_out", ): np.float32, ("weight", "alpha_factor"): np.float_, } domain_range = [(1, 15), (1, 15), (1, 15)] backends = ["debug", "numpy", "gtx86"] symbols = dict( USE_ALPHA=gt_testing.global_name(one_of=(True, False)), field_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_c=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_out=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), weight=gt_testing.parameter(in_range=(-10, 10)), alpha_factor=gt_testing.parameter(in_range=(-1, 1)), ) def definition(field_a, field_b, field_c, field_out, *, weight, alpha_factor): from __externals__ import USE_ALPHA with computation(PARALLEL), interval(...): factor = field_a # TODO remove once scalar assignment issue is fixed if USE_ALPHA: factor = alpha_factor else: factor = 1.0 + 0.0 * ( alpha_factor) # ?? won't compile for just factor = 0.0 field_out = factor * field_a[0, 0, 0] - (1 - factor) * ( field_b[0, 0, 0] - weight * field_c[0, 0, 0]) def validation(field_a, field_b, field_c, field_out, *, weight, alpha_factor, domain, origin, **kwargs): if USE_ALPHA: factor = alpha_factor else: factor = 1.0 field_out[...] = (factor * field_a[:, :, :] ) - (1 - factor) * (field_b[:, :, :] - (weight * field_c[:, :, :]))
class TestParametricMix(gt_testing.StencilTestSuite): """Linear combination of input fields using several parameters.""" dtypes = { ("USE_ALPHA", ): np.int_, ("field_a", "field_b", "field_c"): np.float64, ("field_out", ): np.float32, ("weight", "alpha_factor"): np.float_, } domain_range = [(1, 15), (1, 15), (1, 15)] backends = INTERNAL_BACKENDS symbols = dict( USE_ALPHA=gt_testing.global_name(one_of=(True, False)), field_a=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_b=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_c=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), field_out=gt_testing.field(in_range=(-10, 10), boundary=[(0, 0), (0, 0), (0, 0)]), weight=gt_testing.parameter(in_range=(-10, 10)), alpha_factor=gt_testing.parameter(in_range=(-1, 1)), ) def definition(field_a, field_b, field_c, field_out, *, weight, alpha_factor): from __externals__ import USE_ALPHA from __gtscript__ import __INLINED with computation(PARALLEL), interval(...): if __INLINED(USE_ALPHA): factor = alpha_factor else: factor = 1.0 field_out = factor * field_a[ # noqa: F841 # Local name is assigned to but never used 0, 0, 0] - (1 - factor) * (field_b[0, 0, 0] - weight * field_c[0, 0, 0]) def validation(field_a, field_b, field_c, field_out, *, weight, alpha_factor, domain, origin, **kwargs): if USE_ALPHA: # noqa: F821 # Undefined name factor = alpha_factor else: factor = 1.0 field_out[...] = (factor * field_a[:, :, :] ) - (1 - factor) * (field_b[:, :, :] - (weight * field_c[:, :, :]))
class TestGlobalScale(gt_testing.StencilTestSuite): """Scale stencil using a global global_name.""" dtypes = (np.float_,) domain_range = [(1, 15), (1, 15), (1, 15)] backends = CPU_BACKENDS symbols = dict( SCALE_FACTOR=gt_testing.global_name(one_of=(1.0, 1e3, 1e6)), field_a=gt_testing.field(in_range=(-1, 1), boundary=[(0, 0), (0, 0), (0, 0)]), ) def definition(field_a): from __externals__ import SCALE_FACTOR with computation(PARALLEL), interval(...): field_a = SCALE_FACTOR * field_a[0, 0, 0] def validation(field_a, domain, origin, **kwargs): field_a[...] = SCALE_FACTOR * field_a
class TestNotSpecifiedTwoOptionalFields(TestTwoOptionalFields): backends = INTERNAL_BACKENDS symbols = TestTwoOptionalFields.symbols.copy() symbols["PHYS_TEND_A"] = gt_testing.global_name(one_of=(False, )) symbols["phys_tend_a"] = gt_testing.none()
class TestNotSpecifiedOptionalField(TestOptionalField): backends = INTERNAL_BACKENDS symbols = TestOptionalField.symbols.copy() symbols["PHYS_TEND"] = gt_testing.global_name(one_of=(False, )) symbols["phys_tend"] = gt_testing.none()
def test_gt4py_setup(): fwd_diff = gt_testing.global_name(singleton=fwd_diff_op_xy) assert fwd_diff.kind == "singleton"