From 709dc18550a0d2440a661fc9cef06ce8526f8650 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 3 Apr 2023 12:57:38 +0200 Subject: [PATCH 01/14] First try implementing stencil 20 --- .../mo_solve_nonhydro_stencil_20.py | 208 ++++-------------- 1 file changed, 43 insertions(+), 165 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index 06a3829b29..fd88ad415a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -11,176 +11,54 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - list_get, - named_range, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.experimental import as_offset -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.icon4pygen.metadata import FieldInfo +from icon4py.common.dimension import E2C, CellDim, E2EC, ECDim, EdgeDim, KDim, Koff - -@fundef -def step( - i, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, -): - d_ikoffset = list_get(i, deref(ikoffset)) - d_z_exner_exp_pr = deref(shift(Koff, d_ikoffset, E2C, i)(z_exner_ex_pr)) - d_z_dexner_dz_c_1 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_1)) - d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) - d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) - - return d_z_exner_exp_pr + d_zdiff_gradp * ( - d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 - ) - - -@fundef +@field_operator def _mo_solve_nonhydro_stencil_20( - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, -): - return deref(inv_dual_edge_length) * ( - step(1, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2) - - step( - 0, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2 - ) - ) + inv_dual_edge_length: Field[[EdgeDim], float], + z_exner_ex_pr: Field[[CellDim, KDim], float], + zdiff_gradp: Field[[EdgeDim, ECDim, KDim], float], + ikoffset: Field[[EdgeDim, ECDim, KDim], int], + z_dexner_dz_c_1: Field[[CellDim, KDim], float], + z_dexner_dz_c_2: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + #z_gradh_exner = inv_dual_edge_length * (( + #z_exner_ex_pr(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) + + #zdiff_gradp(E2EC[1]) * ( + #z_dexner_dz_c_1(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) + + #z_dexner_dz_c_2(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) + #)) - ( + #z_exner_ex_pr(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + + #zdiff_gradp(E2EC[0]) * ( + #z_dexner_dz_c_1(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + + #z_dexner_dz_c_2(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + #))) + + temp = z_exner_ex_pr(E2C[1]) + z_gradh_exner = temp(as_offset(Koff, ikoffset(E2EC[1]))) + return z_gradh_exner -@fendef +@program def mo_solve_nonhydro_stencil_20( - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, - z_gradh_exner, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, + inv_dual_edge_length: Field[[EdgeDim], float], + z_exner_ex_pr: Field[[CellDim, KDim], float], + zdiff_gradp: Field[[EdgeDim, ECDim, KDim], float], + ikoffset: Field[[EdgeDim, ECDim, KDim], int], + z_dexner_dz_c_1: Field[[CellDim, KDim], float], + z_dexner_dz_c_2: Field[[CellDim, KDim], float], + z_gradh_exner: Field[[EdgeDim, KDim], float], ): - closure( - unstructured_domain( - named_range(EdgeDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_solve_nonhydro_stencil_20, - z_gradh_exner, - [ - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, - ], + _mo_solve_nonhydro_stencil_20( + inv_dual_edge_length, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, + out=z_gradh_exner, ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "inv_dual_edge_length": FieldInfo( - field=past.FieldSymbol( - id="inv_dual_edge_length", - type=ts.FieldType( - dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_exner_ex_pr": FieldInfo( - field=past.FieldSymbol( - id="z_exner_ex_pr", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zdiff_gradp": FieldInfo( - field=past.FieldSymbol( - id="zdiff_gradp", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "ikoffset": FieldInfo( - field=past.FieldSymbol( - id="ikoffset", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_dexner_dz_c_1": FieldInfo( - field=past.FieldSymbol( - id="z_dexner_dz_c_1", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_dexner_dz_c_2": FieldInfo( - field=past.FieldSymbol( - id="z_dexner_dz_c_2", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_gradh_exner": FieldInfo( - field=past.FieldSymbol( - id="z_gradh_exner", - type=ts.FieldType( - dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=False, - out=True, - ), -} -# patch the fendef with metainfo for icon4pygen -mo_solve_nonhydro_stencil_20.__dict__["offsets"] = [ - Koff.value, - E2C.value, -] # could be done with a pass... -mo_solve_nonhydro_stencil_20.__dict__["metadata"] = _metadata From b3cf80c1808d6f9809d3dab2bc23b7924cb03da8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 5 Apr 2023 18:15:44 +0200 Subject: [PATCH 02/14] Completed stencil, wrote test --- .../mo_solve_nonhydro_stencil_20.py | 46 ++++++++++++------- ...ate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../test_mo_solve_nonhydro_stencil_20.py | 16 +++++-- testutils/src/icon4py/testutils/utils.py | 40 ++++++++++++++++ 4 files changed, 82 insertions(+), 22 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index fd88ad415a..56e36ccf3e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -21,25 +21,37 @@ def _mo_solve_nonhydro_stencil_20( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], - zdiff_gradp: Field[[EdgeDim, ECDim, KDim], float], - ikoffset: Field[[EdgeDim, ECDim, KDim], int], + zdiff_gradp: Field[[ECDim, KDim], float], + ikoffset: Field[[ECDim, KDim], int], z_dexner_dz_c_1: Field[[CellDim, KDim], float], z_dexner_dz_c_2: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - #z_gradh_exner = inv_dual_edge_length * (( - #z_exner_ex_pr(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) + - #zdiff_gradp(E2EC[1]) * ( - #z_dexner_dz_c_1(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) + - #z_dexner_dz_c_2(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) - #)) - ( - #z_exner_ex_pr(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + - #zdiff_gradp(E2EC[0]) * ( - #z_dexner_dz_c_1(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + - #z_dexner_dz_c_2(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) - #))) - temp = z_exner_ex_pr(E2C[1]) - z_gradh_exner = temp(as_offset(Koff, ikoffset(E2EC[1]))) + z_exner_ex_pr_0 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[0]))) + z_exner_ex_pr_1 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[1]))) + + z_dexner_dz_c1_0 = z_dexner_dz_c_1(as_offset(Koff, ikoffset(E2EC[0]))) + z_dexner_dz_c1_1 = z_dexner_dz_c_1(as_offset(Koff, ikoffset(E2EC[1]))) + + z_dexner_dz_c2_0 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[0]))) + z_dexner_dz_c2_1 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[1]))) + + z_gradh_exner = inv_dual_edge_length * (( + z_exner_ex_pr_1(E2C[1]) + + zdiff_gradp(E2EC[1]) * ( + z_dexner_dz_c1_1(E2C[1]) + + z_dexner_dz_c2_1(E2C[1]) + )) - ( + z_exner_ex_pr_0(E2C[0]) + + zdiff_gradp(E2EC[0]) * ( + z_dexner_dz_c1_0(E2C[0]) + + z_dexner_dz_c2_0(E2C[0]) + ))) + + #z_exner_ex_pr_0 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[0]))) + + #z_gradh_exner = z_exner_ex_pr_0(E2C[0]) + return z_gradh_exner @@ -47,8 +59,8 @@ def _mo_solve_nonhydro_stencil_20( def mo_solve_nonhydro_stencil_20( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], - zdiff_gradp: Field[[EdgeDim, ECDim, KDim], float], - ikoffset: Field[[EdgeDim, ECDim, KDim], int], + zdiff_gradp: Field[[ECDim, KDim], float], + ikoffset: Field[[ECDim, KDim], int], z_dexner_dz_c_1: Field[[CellDim, KDim], float], z_dexner_dz_c_2: Field[[CellDim, KDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 967b2b1532..c73210845d 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -213,6 +213,6 @@ def test_calculate_nabla2_and_smag_coefficients_for_vn(): "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, mesh.n_e2c2v), }, ) - assert np.allclose(kh_smag_e_ref, kh_smag_e) + assert not np.allclose(kh_smag_e_ref, kh_smag_e) assert np.allclose(kh_smag_ec_ref, kh_smag_ec) assert np.allclose(z_nabla2_e_ref, z_nabla2_e) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 040a88ba25..1951881e9d 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -12,13 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, ) -from icon4py.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field def mo_solve_nonhydro_stencil_20_numpy( @@ -71,6 +72,7 @@ def test_mo_solve_nonhydro_stencil_20(): z_exner_ex_pr = random_field(mesh, CellDim, KDim) zdiff_gradp = random_field(mesh, EdgeDim, E2CDim, KDim) ikoffset = zero_field(mesh, EdgeDim, E2CDim, KDim, dtype=int) + rng = np.random.default_rng() for k in range(mesh.k_level): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -80,6 +82,9 @@ def test_mo_solve_nonhydro_stencil_20(): size=(ikoffset.shape[0], ikoffset.shape[1]), ) + zdiff_gradp_new = flatten_first_two_dims(ECDim, KDim, field=zdiff_gradp) + ikoffset_new = flatten_first_two_dims(ECDim, KDim, field=ikoffset) + z_dexner_dz_c_1 = random_field(mesh, CellDim, KDim) z_dexner_dz_c_2 = random_field(mesh, CellDim, KDim) @@ -103,8 +108,8 @@ def test_mo_solve_nonhydro_stencil_20(): mo_solve_nonhydro_stencil_20( inv_dual_edge_length, z_exner_ex_pr, - zdiff_gradp, - ikoffset, + zdiff_gradp_new, + ikoffset_new, z_dexner_dz_c_1, z_dexner_dz_c_2, z_gradh_exner, @@ -114,8 +119,11 @@ def test_mo_solve_nonhydro_stencil_20(): kend, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), "Koff": KDim, }, ) + print(np.asarray(z_gradh_exner)) + print(z_gradh_exner_ref) assert np.allclose(z_gradh_exner, z_gradh_exner_ref) diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index e4d503cfca..428e7edbba 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -87,11 +87,51 @@ def as_1D_sparse_field( field: it_embedded.MutableLocatedField, dim: gt_common.Dimension ) -> it_embedded.MutableLocatedField: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" + print("OldArray") + print(np.asarray(field)) old_shape = np.asarray(field).shape + print(old_shape) assert len(old_shape) == 2 new_shape = (old_shape[0] * old_shape[1],) + print("NEWArray") + print(np.asarray(field).reshape(new_shape)) + print(new_shape) return it_embedded.np_as_located_field(dim)(np.asarray(field).reshape(new_shape)) +def flatten_first_two_dims( + *dims: gt_common.Dimension, field: it_embedded.MutableLocatedField +) -> it_embedded.MutableLocatedField: + """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" + old_shape = np.asarray(field).shape + print("Old Shape") + print(old_shape) + assert len(old_shape) >= 2 + flattened_size = old_shape[0] * old_shape[1] + print("Flattened Size") + print(flattened_size) + flattened_shape = (flattened_size, ) + print("Flattened Shape") + print(flattened_shape) + residual_shape = old_shape[2:] + print("Residual Shape") + print(residual_shape) + new_shape = flattened_shape + residual_shape + print("New Shape") + print(new_shape) + + print("Field") + print(field) + array = np.asarray(field) + print("ARRAY") + print(array) + print(array.shape) + newarray = np.asarray(field).reshape(new_shape) + print("NEWARRAY") + print(newarray) + print(newarray.shape) + print(dims) + return it_embedded.np_as_located_field(*dims)(newarray) + def get_stencil_module_path(stencil_module: str, stencil_name: str) -> str: return f"icon4py.{stencil_module}.{stencil_name}:{stencil_name}" From 6f07be7de3fdbe332d1f548bb592a41015c357c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 21 Apr 2023 16:41:34 +0200 Subject: [PATCH 03/14] Re-introduce 3rd ITIR version for verification --- .../mo_solve_nonhydro_stencil_20_itir.py | 186 ++++++++++++++++++ ...ate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../test_mo_solve_nonhydro_stencil_20.py | 30 ++- 3 files changed, 215 insertions(+), 3 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py new file mode 100644 index 0000000000..2bcd63c5c1 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py @@ -0,0 +1,186 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.eve import SourceLocation +from gt4py.next.ffront import program_ast as past +from gt4py.next.iterator.builtins import ( + deref, + list_get, + named_range, + shift, + unstructured_domain, +) +from gt4py.next.iterator.runtime import closure, fendef, fundef +from gt4py.next.type_system import type_specifications as ts + +from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff +from icon4py.pyutils.metadata import FieldInfo + + +@fundef +def step( + i, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, +): + d_ikoffset = list_get(i, deref(ikoffset)) + d_z_exner_exp_pr = deref(shift(Koff, d_ikoffset, E2C, i)(z_exner_ex_pr)) + d_z_dexner_dz_c_1 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_1)) + d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) + d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) + + return d_z_exner_exp_pr + d_zdiff_gradp * ( + d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 + ) + + +@fundef +def _mo_solve_nonhydro_stencil_20_itir( + inv_dual_edge_length, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, +): + return deref(inv_dual_edge_length) * ( + step(1, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2) + - step( + 0, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2 + ) + ) + + +@fendef +def mo_solve_nonhydro_stencil_20_itir( + inv_dual_edge_length, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, + z_gradh_exner, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, +): + closure( + unstructured_domain( + named_range(EdgeDim, horizontal_start, horizontal_end), + named_range(KDim, vertical_start, vertical_end), + ), + _mo_solve_nonhydro_stencil_20_itir, + z_gradh_exner, + [ + inv_dual_edge_length, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, + ], + ) + + +_dummy_loc = SourceLocation(1, 1, "") +_metadata = { + "inv_dual_edge_length": FieldInfo( + field=past.FieldSymbol( + id="inv_dual_edge_length", + type=ts.FieldType( + dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_exner_ex_pr": FieldInfo( + field=past.FieldSymbol( + id="z_exner_ex_pr", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "zdiff_gradp": FieldInfo( + field=past.FieldSymbol( + id="zdiff_gradp", + type=ts.FieldType( + dims=[EdgeDim, E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "ikoffset": FieldInfo( + field=past.FieldSymbol( + id="ikoffset", + type=ts.FieldType( + dims=[EdgeDim, E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_dexner_dz_c_1": FieldInfo( + field=past.FieldSymbol( + id="z_dexner_dz_c_1", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_dexner_dz_c_2": FieldInfo( + field=past.FieldSymbol( + id="z_dexner_dz_c_2", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_gradh_exner": FieldInfo( + field=past.FieldSymbol( + id="z_gradh_exner", + type=ts.FieldType( + dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=False, + out=True, + ), +} +# patch the fendef with metainfo for icon4pygen +mo_solve_nonhydro_stencil_20_itir.__dict__["offsets"] = [ + Koff.value, + E2C.value, +] # could be done with a pass... +mo_solve_nonhydro_stencil_20_itir.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index c73210845d..967b2b1532 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -213,6 +213,6 @@ def test_calculate_nabla2_and_smag_coefficients_for_vn(): "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, mesh.n_e2c2v), }, ) - assert not np.allclose(kh_smag_e_ref, kh_smag_e) + assert np.allclose(kh_smag_e_ref, kh_smag_e) assert np.allclose(kh_smag_ec_ref, kh_smag_ec) assert np.allclose(z_nabla2_e_ref, z_nabla2_e) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 1951881e9d..69ace485b4 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -17,6 +17,11 @@ from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, ) + +from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20_itir import ( + mo_solve_nonhydro_stencil_20_itir, +) + from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field @@ -89,6 +94,7 @@ def test_mo_solve_nonhydro_stencil_20(): z_dexner_dz_c_2 = random_field(mesh, CellDim, KDim) z_gradh_exner = zero_field(mesh, EdgeDim, KDim) + z_gradh_exner_itir = zero_field(mesh, EdgeDim, KDim) z_gradh_exner_ref = mo_solve_nonhydro_stencil_20_numpy( mesh.e2c, @@ -105,6 +111,24 @@ def test_mo_solve_nonhydro_stencil_20(): kstart = 0 kend = mesh.k_level + mo_solve_nonhydro_stencil_20_itir( + inv_dual_edge_length, + z_exner_ex_pr, + zdiff_gradp, + ikoffset, + z_dexner_dz_c_1, + z_dexner_dz_c_2, + z_gradh_exner_itir, + hstart, + hend, + kstart, + kend, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "Koff": KDim, + }, + ) + mo_solve_nonhydro_stencil_20( inv_dual_edge_length, z_exner_ex_pr, @@ -124,6 +148,8 @@ def test_mo_solve_nonhydro_stencil_20(): }, ) - print(np.asarray(z_gradh_exner)) - print(z_gradh_exner_ref) + #print(np.asarray(z_gradh_exner)) + #print(z_gradh_exner_ref) + + assert np.allclose(z_gradh_exner_itir, z_gradh_exner_ref) assert np.allclose(z_gradh_exner, z_gradh_exner_ref) From aa01bcaa4713f64bd434b5071db936c3793a44e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 21 Apr 2023 17:49:35 +0200 Subject: [PATCH 04/14] First passing version --- .../mo_solve_nonhydro_stencil_20.py | 2 +- .../mo_solve_nonhydro_stencil_20_itir.py | 4 +- .../test_mo_solve_nonhydro_stencil_20.py | 24 +++++---- testutils/src/icon4py/testutils/utils.py | 50 +++++++++---------- 4 files changed, 40 insertions(+), 40 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index 56e36ccf3e..67de4532e2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -52,7 +52,7 @@ def _mo_solve_nonhydro_stencil_20( #z_gradh_exner = z_exner_ex_pr_0(E2C[0]) - return z_gradh_exner + return inv_dual_edge_length * (z_exner_ex_pr_1(E2C[1]) - z_exner_ex_pr_0(E2C[0])) @program diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py index 2bcd63c5c1..4a452a2719 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py @@ -42,9 +42,7 @@ def step( d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) - return d_z_exner_exp_pr + d_zdiff_gradp * ( - d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 - ) + return d_z_exner_exp_pr @fundef diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 69ace485b4..35124a2679 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -51,18 +51,19 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) z_exner_ex_pr_at_kidx = _apply_index_field(full_shape, z_exner_ex_pr, e2c, ikoffset) - z_dexner_dz_c_1_at_kidx = _apply_index_field( - full_shape, z_dexner_dz_c_1, e2c, ikoffset - ) - z_dexner_dz_c_2_at_kidx = _apply_index_field( - full_shape, z_dexner_dz_c_2, e2c, ikoffset - ) + #z_dexner_dz_c_1_at_kidx = _apply_index_field( + # full_shape, z_dexner_dz_c_1, e2c, ikoffset + #) + #z_dexner_dz_c_2_at_kidx = _apply_index_field( + # full_shape, z_dexner_dz_c_2, e2c, ikoffset + #) def at_neighbor(i): - return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( - z_dexner_dz_c_1_at_kidx[:, i, :] - + zdiff_gradp[:, i, :] * z_dexner_dz_c_2_at_kidx[:, i, :] - ) + return z_exner_ex_pr_at_kidx[:, i, :] + #return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( + # z_dexner_dz_c_1_at_kidx[:, i, :] + # + zdiff_gradp[:, i, :] * z_dexner_dz_c_2_at_kidx[:, i, :] + #) sum_expr = at_neighbor(1) - at_neighbor(0) @@ -148,8 +149,9 @@ def test_mo_solve_nonhydro_stencil_20(): }, ) - #print(np.asarray(z_gradh_exner)) #print(z_gradh_exner_ref) + #print(z_gradh_exner_itir ) + #print(np.asarray(z_gradh_exner)) assert np.allclose(z_gradh_exner_itir, z_gradh_exner_ref) assert np.allclose(z_gradh_exner, z_gradh_exner_ref) diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index 428e7edbba..eb4e497d3f 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -87,15 +87,15 @@ def as_1D_sparse_field( field: it_embedded.MutableLocatedField, dim: gt_common.Dimension ) -> it_embedded.MutableLocatedField: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" - print("OldArray") - print(np.asarray(field)) + #print("OldArray") + #print(np.asarray(field)) old_shape = np.asarray(field).shape - print(old_shape) + #print(old_shape) assert len(old_shape) == 2 new_shape = (old_shape[0] * old_shape[1],) - print("NEWArray") - print(np.asarray(field).reshape(new_shape)) - print(new_shape) + #print("NEWArray") + #print(np.asarray(field).reshape(new_shape)) + #print(new_shape) return it_embedded.np_as_located_field(dim)(np.asarray(field).reshape(new_shape)) def flatten_first_two_dims( @@ -103,33 +103,33 @@ def flatten_first_two_dims( ) -> it_embedded.MutableLocatedField: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" old_shape = np.asarray(field).shape - print("Old Shape") - print(old_shape) + #print("Old Shape") + #print(old_shape) assert len(old_shape) >= 2 flattened_size = old_shape[0] * old_shape[1] - print("Flattened Size") - print(flattened_size) + #print("Flattened Size") + #print(flattened_size) flattened_shape = (flattened_size, ) - print("Flattened Shape") - print(flattened_shape) + #print("Flattened Shape") + #print(flattened_shape) residual_shape = old_shape[2:] - print("Residual Shape") - print(residual_shape) + #print("Residual Shape") + #print(residual_shape) new_shape = flattened_shape + residual_shape - print("New Shape") - print(new_shape) + #print("New Shape") + #print(new_shape) - print("Field") - print(field) + #print("Field") + #print(field) array = np.asarray(field) - print("ARRAY") - print(array) - print(array.shape) + #print("ARRAY") + #print(array) + #print(array.shape) newarray = np.asarray(field).reshape(new_shape) - print("NEWARRAY") - print(newarray) - print(newarray.shape) - print(dims) + #print("NEWARRAY") + #print(newarray) + #print(newarray.shape) + #print(dims) return it_embedded.np_as_located_field(*dims)(newarray) From d9ff32696032fe57a29aed445109e517a39b07d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 26 Apr 2023 10:12:41 +0200 Subject: [PATCH 05/14] stencil 15 diffusion working --- .../mo_nh_diffusion_stencil_15.py | 217 ++++-------------- .../mo_nh_diffusion_stencil_15_itir.py | 198 ++++++++++++++++ .../mo_solve_nonhydro_stencil_20.py | 10 +- .../mo_solve_nonhydro_stencil_20_itir.py | 4 +- .../tests/test_mo_nh_diffusion_stencil_15.py | 19 +- .../test_mo_solve_nonhydro_stencil_20.py | 21 +- 6 files changed, 272 insertions(+), 197 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index a0d3fca5b1..63116f3769 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -11,187 +11,58 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - if_, - list_get, - named_range, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts - -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, Koff -from icon4py.icon4pygen.metadata import FieldInfo +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, where +from gt4py.next.ffront.experimental import as_offset +from icon4py.common.dimension import C2E2C, CellDim, CECDim, C2CEC, KDim, Koff +@field_operator +def _mo_nh_diffusion_stencil_15( + mask: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim, KDim], int], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s_c: Field[[CellDim], float], + geofac_n2s_nbh: Field[[CECDim], float], + vcoef: Field[[CECDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + z_temp: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: -def step(i, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset): - d_vcoef = list_get(i, deref(vcoef)) - s_theta_v = shift(C2E2C, i, Koff, list_get(i, deref(zd_vertoffset)))(theta_v) - return list_get(i, deref(geofac_n2s_nbh)) * ( - d_vcoef * deref(s_theta_v) + (1.0 - d_vcoef) * deref(shift(Koff, 1)(s_theta_v)) - ) + theta_v_0 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]))) + theta_v_1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]))) + theta_v_2 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]))) + theta_v_0_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]) + 1)) + theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + 1)) + theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + 1)) -@fundef -def _mo_nh_diffusion_stencil_15( - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, -): - summed = ( - step(0, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - + step(1, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - + step(2, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - ) - update = deref(z_temp) + deref(zd_diffcoef) * ( - deref(theta_v) * deref(geofac_n2s_c) + summed - ) + sum = geofac_n2s_nbh(C2CEC[0]) * (vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0])) \ + + geofac_n2s_nbh(C2CEC[1]) * (vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1])) \ + + geofac_n2s_nbh(C2CEC[2]) * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) - return if_(deref(mask), update, deref(z_temp)) + z_temp = where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum), z_temp) + return z_temp -@fendef +@program def mo_nh_diffusion_stencil_15( - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, + mask: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim, KDim], int], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s_c: Field[[CellDim], float], + geofac_n2s_nbh: Field[[CECDim], float], + vcoef: Field[[CECDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + z_temp: Field[[CellDim, KDim], float], ): - closure( - unstructured_domain( - named_range(CellDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_nh_diffusion_stencil_15, + _mo_nh_diffusion_stencil_15( + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v, z_temp, - [ - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, - ], + out=z_temp, ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "mask": FieldInfo( - field=past.FieldSymbol( - id="mask", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zd_vertoffset": FieldInfo( - field=past.FieldSymbol( - id="zd_vertoffset", - type=ts.FieldType( - dims=[CellDim, C2E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zd_diffcoef": FieldInfo( - field=past.FieldSymbol( - id="zd_diffcoef", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "geofac_n2s_c": FieldInfo( - field=past.FieldSymbol( - id="geofac_n2s_c", - type=ts.FieldType( - dims=[CellDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "geofac_n2s_nbh": FieldInfo( - field=past.FieldSymbol( - id="geofac_n2s_nbh", - type=ts.FieldType( - dims=[CellDim, C2E2CDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "vcoef": FieldInfo( - field=past.FieldSymbol( - id="vcoef", - type=ts.FieldType( - dims=[CellDim, C2E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "theta_v": FieldInfo( - field=past.FieldSymbol( - id="theta_v", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_temp": FieldInfo( - field=past.FieldSymbol( - id="z_temp", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=True, - ), -} - -# patch the fendef with metainfo for icon4pygen -mo_nh_diffusion_stencil_15.__dict__["offsets"] = [ - Koff.value, - C2E2C.value, -] # could be done with a pass... -mo_nh_diffusion_stencil_15.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py new file mode 100644 index 0000000000..db4a421745 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py @@ -0,0 +1,198 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.eve import SourceLocation +from gt4py.next.ffront import program_ast as past +from gt4py.next.iterator.builtins import ( + deref, + if_, + list_get, + named_range, + shift, + unstructured_domain, +) +from gt4py.next.iterator.runtime import closure, fendef, fundef +from gt4py.next.type_system import type_specifications as ts + +from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, Koff +from icon4py.pyutils.metadata import FieldInfo + + +@fundef +def step(i, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset): + d_vcoef = list_get(i, deref(vcoef)) + s_theta_v = shift(C2E2C, i, Koff, list_get(i, deref(zd_vertoffset)))(theta_v) + return list_get(i, deref(geofac_n2s_nbh)) * ( + d_vcoef * deref(s_theta_v) + (1.0 - d_vcoef) * deref(shift(Koff, 1)(s_theta_v)) + ) + + +@fundef +def _mo_nh_diffusion_stencil_15_itir( + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v, + z_temp, +): + summed = ( + step(0, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) + + step(1, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) + + step(2, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) + ) + update = deref(z_temp) + deref(zd_diffcoef) * ( + deref(theta_v) * deref(geofac_n2s_c) + summed + ) + + return if_(deref(mask), update, deref(z_temp)) + + +@fendef +def mo_nh_diffusion_stencil_15_itir( + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v, + z_temp, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, +): + closure( + unstructured_domain( + named_range(CellDim, horizontal_start, horizontal_end), + named_range(KDim, vertical_start, vertical_end), + ), + _mo_nh_diffusion_stencil_15_itir, + z_temp, + [ + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v, + z_temp, + ], + ) + + +_dummy_loc = SourceLocation(1, 1, "") +_metadata = { + "mask": FieldInfo( + field=past.FieldSymbol( + id="mask", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "zd_vertoffset": FieldInfo( + field=past.FieldSymbol( + id="zd_vertoffset", + type=ts.FieldType( + dims=[CellDim, C2E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "zd_diffcoef": FieldInfo( + field=past.FieldSymbol( + id="zd_diffcoef", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "geofac_n2s_c": FieldInfo( + field=past.FieldSymbol( + id="geofac_n2s_c", + type=ts.FieldType( + dims=[CellDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "geofac_n2s_nbh": FieldInfo( + field=past.FieldSymbol( + id="geofac_n2s_nbh", + type=ts.FieldType( + dims=[CellDim, C2E2CDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "vcoef": FieldInfo( + field=past.FieldSymbol( + id="vcoef", + type=ts.FieldType( + dims=[CellDim, C2E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "theta_v": FieldInfo( + field=past.FieldSymbol( + id="theta_v", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_temp": FieldInfo( + field=past.FieldSymbol( + id="z_temp", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=True, + ), +} + +# patch the fendef with metainfo for icon4pygen +mo_nh_diffusion_stencil_15_itir.__dict__["offsets"] = [ + Koff.value, + C2E2C.value, +] # could be done with a pass... +mo_nh_diffusion_stencil_15_itir.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index 67de4532e2..a96074291c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -40,19 +40,15 @@ def _mo_solve_nonhydro_stencil_20( z_exner_ex_pr_1(E2C[1]) + zdiff_gradp(E2EC[1]) * ( z_dexner_dz_c1_1(E2C[1]) + - z_dexner_dz_c2_1(E2C[1]) + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1]) )) - ( z_exner_ex_pr_0(E2C[0]) + zdiff_gradp(E2EC[0]) * ( z_dexner_dz_c1_0(E2C[0]) + - z_dexner_dz_c2_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) ))) - #z_exner_ex_pr_0 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[0]))) - - #z_gradh_exner = z_exner_ex_pr_0(E2C[0]) - - return inv_dual_edge_length * (z_exner_ex_pr_1(E2C[1]) - z_exner_ex_pr_0(E2C[0])) + return z_gradh_exner @program diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py index 4a452a2719..2bcd63c5c1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py @@ -42,7 +42,9 @@ def step( d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) - return d_z_exner_exp_pr + return d_z_exner_exp_pr + d_zdiff_gradp * ( + d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 + ) @fundef diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 829fb9391c..1fc2d65383 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -13,12 +13,16 @@ import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( mo_nh_diffusion_stencil_15, ) -from icon4py.common.dimension import C2E2CDim, CellDim, KDim +from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, random_mask, zero_field +from icon4py.testutils.utils import flatten_first_two_dims, random_field, random_mask, zero_field + +from gt4py.next.iterator import embedded as it_embedded def mo_nh_diffusion_stencil_15_numpy( @@ -82,6 +86,10 @@ def test_mo_nh_diffusion_stencil_15(): theta_v = random_field(mesh, CellDim, KDim) z_temp = random_field(mesh, CellDim, KDim) + vcoef_new = flatten_first_two_dims(CECDim, KDim, field=vcoef) + zd_vertoffset_new = flatten_first_two_dims(CECDim, KDim, field=zd_vertoffset) + geofac_n2s_nbh_new = flatten_first_two_dims(CECDim, field=geofac_n2s_nbh) + ref = mo_nh_diffusion_stencil_15_numpy( mesh.c2e2c, np.asarray(mask), @@ -101,11 +109,11 @@ def test_mo_nh_diffusion_stencil_15(): mo_nh_diffusion_stencil_15( mask, - zd_vertoffset, + zd_vertoffset_new, zd_diffcoef, geofac_n2s_c, - geofac_n2s_nbh, - vcoef, + geofac_n2s_nbh_new, + vcoef_new, theta_v, z_temp, hstart, @@ -114,6 +122,7 @@ def test_mo_nh_diffusion_stencil_15(): kend, offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), + "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), "Koff": KDim, }, ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 35124a2679..bf10a8ea32 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -51,19 +51,18 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) z_exner_ex_pr_at_kidx = _apply_index_field(full_shape, z_exner_ex_pr, e2c, ikoffset) - #z_dexner_dz_c_1_at_kidx = _apply_index_field( - # full_shape, z_dexner_dz_c_1, e2c, ikoffset - #) - #z_dexner_dz_c_2_at_kidx = _apply_index_field( - # full_shape, z_dexner_dz_c_2, e2c, ikoffset - #) + z_dexner_dz_c_1_at_kidx = _apply_index_field( + full_shape, z_dexner_dz_c_1, e2c, ikoffset + ) + z_dexner_dz_c_2_at_kidx = _apply_index_field( + full_shape, z_dexner_dz_c_2, e2c, ikoffset + ) def at_neighbor(i): - return z_exner_ex_pr_at_kidx[:, i, :] - #return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( - # z_dexner_dz_c_1_at_kidx[:, i, :] - # + zdiff_gradp[:, i, :] * z_dexner_dz_c_2_at_kidx[:, i, :] - #) + return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( + z_dexner_dz_c_1_at_kidx[:, i, :] + + zdiff_gradp[:, i, :] * z_dexner_dz_c_2_at_kidx[:, i, :] + ) sum_expr = at_neighbor(1) - at_neighbor(0) From d2718c8a9d5979e38ca827afd6f762fbd25b0922 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 2 May 2023 00:45:11 +0200 Subject: [PATCH 06/14] Fixed bindings generator, ported rest of the stencils --- .../mo_nh_diffusion_stencil_15.py | 12 +- .../atm_dyn_iconam/mo_solve_nonhydro_20.py | 0 .../mo_solve_nonhydro_stencil_20.py | 6 +- .../mo_solve_nonhydro_stencil_21.py | 219 ++++-------------- .../mo_solve_nonhydro_stencil_21_itir.py | 200 ++++++++++++++++ .../test_mo_solve_nonhydro_stencil_20.py | 3 +- .../test_mo_solve_nonhydro_stencil_21.py | 18 +- .../bindings/codegen/render/field.py | 18 +- .../icon4py/icon4pygen/bindings/entities.py | 4 +- .../icon4py/icon4pygen/bindings/locations.py | 9 + 10 files changed, 288 insertions(+), 201 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 63116f3769..71279d98b8 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -12,14 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, where +from gt4py.next.ffront.fbuiltins import Field, where, int32 from gt4py.next.ffront.experimental import as_offset from icon4py.common.dimension import C2E2C, CellDim, CECDim, C2CEC, KDim, Koff @field_operator def _mo_nh_diffusion_stencil_15( mask: Field[[CellDim, KDim], bool], - zd_vertoffset: Field[[CECDim, KDim], int], + zd_vertoffset: Field[[CECDim, KDim], int32], zd_diffcoef: Field[[CellDim, KDim], float], geofac_n2s_c: Field[[CellDim], float], geofac_n2s_nbh: Field[[CECDim], float], @@ -32,9 +32,9 @@ def _mo_nh_diffusion_stencil_15( theta_v_1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]))) theta_v_2 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]))) - theta_v_0_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]) + 1)) - theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + 1)) - theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + 1)) + theta_v_0_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]) + int32(1))) + theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + int32(1))) + theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + int32(1))) sum = geofac_n2s_nbh(C2CEC[0]) * (vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0])) \ + geofac_n2s_nbh(C2CEC[1]) * (vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1])) \ @@ -47,7 +47,7 @@ def _mo_nh_diffusion_stencil_15( @program def mo_nh_diffusion_stencil_15( mask: Field[[CellDim, KDim], bool], - zd_vertoffset: Field[[CECDim, KDim], int], + zd_vertoffset: Field[[CECDim, KDim], int32], zd_diffcoef: Field[[CellDim, KDim], float], geofac_n2s_c: Field[[CellDim], float], geofac_n2s_nbh: Field[[CECDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index a96074291c..fb3abd55b4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, neighbor_sum, int32 from gt4py.next.ffront.experimental import as_offset from icon4py.common.dimension import E2C, CellDim, E2EC, ECDim, EdgeDim, KDim, Koff @@ -22,7 +22,7 @@ def _mo_solve_nonhydro_stencil_20( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], zdiff_gradp: Field[[ECDim, KDim], float], - ikoffset: Field[[ECDim, KDim], int], + ikoffset: Field[[ECDim, KDim], int32], z_dexner_dz_c_1: Field[[CellDim, KDim], float], z_dexner_dz_c_2: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: @@ -56,7 +56,7 @@ def mo_solve_nonhydro_stencil_20( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], zdiff_gradp: Field[[ECDim, KDim], float], - ikoffset: Field[[ECDim, KDim], int], + ikoffset: Field[[ECDim, KDim], int32], z_dexner_dz_c_1: Field[[CellDim, KDim], float], z_dexner_dz_c_2: Field[[CellDim, KDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py index 4fcd5229a8..f6409ed5c3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py @@ -11,190 +11,61 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - list_get, - named_range, - power, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, neighbor_sum, int32 +from gt4py.next.ffront.experimental import as_offset -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.icon4pygen.metadata import FieldInfo +from icon4py.common.dimension import E2C, CellDim, E2EC, ECDim, EdgeDim, KDim, Koff +@field_operator +def _mo_solve_nonhydro_stencil_21( + theta_v: Field[[CellDim, KDim], float], + ikoffset: Field[[ECDim, KDim], int32], + zdiff_gradp: Field[[ECDim, KDim], float], + theta_v_ic: Field[[CellDim, KDim], float], + inv_ddqz_z_full: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + grav_o_cpd: float, +) -> Field[[EdgeDim, KDim], float]: -@fundef -def step(i, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full): - d_ikoffset = list_get(i, deref(ikoffset)) + theta_v_0 = theta_v(as_offset(Koff, ikoffset(E2EC[0]))) + theta_v_1 = theta_v(as_offset(Koff, ikoffset(E2EC[1]))) - d_theta_v = deref(shift(Koff, d_ikoffset, E2C, i)(theta_v)) - s_theta_v_ic = shift(Koff, d_ikoffset, E2C, i)(theta_v_ic) - d_theta_v_ic = deref(s_theta_v_ic) - d_theta_v_ic_p1 = deref(shift(Koff, 1)(s_theta_v_ic)) - d_inv_ddqz_z_full = deref(shift(Koff, d_ikoffset, E2C, i)(inv_ddqz_z_full)) - d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) + theta_v_ic_0 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[0]))) + theta_v_ic_1 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[1]))) - return ( - d_theta_v + d_zdiff_gradp * (d_theta_v_ic - d_theta_v_ic_p1) * d_inv_ddqz_z_full - ) + theta_v_ic_p1_0 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[0]) + int32(1))) + theta_v_ic_p1_1 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[1]) + int32(1))) + inv_ddqz_z_full_0 = inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[0]))) + inv_ddqz_z_full_1 = inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[1]))) + + z_theta_0 = theta_v_0(E2C[0]) + zdiff_gradp(E2EC[0]) * (theta_v_ic_0(E2C[0]) - theta_v_ic_p1_0(E2C[0])) * inv_ddqz_z_full_0(E2C[0]) + z_theta_1 = theta_v_1(E2C[1]) + zdiff_gradp(E2EC[1]) * (theta_v_ic_1(E2C[1]) - theta_v_ic_p1_1(E2C[1])) * inv_ddqz_z_full_1(E2C[1]) + + z_hydro_corr = grav_o_cpd * inv_dual_edge_length * (z_theta_1 - z_theta_0) * float(4.0) / (z_theta_0 + z_theta_1)**2 -@fundef -def _mo_solve_nonhydro_stencil_21( - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, -): - z_theta1 = step(0, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) - z_theta2 = step(1, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) - z_hydro_corr = ( - deref(grav_o_cpd) - * deref(inv_dual_edge_length) - * (z_theta2 - z_theta1) - * 4.0 - / power((z_theta1 + z_theta2), 2.0) - ) return z_hydro_corr -@fendef +@program def mo_solve_nonhydro_stencil_21( - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, - z_hydro_corr, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, + theta_v: Field[[CellDim, KDim], float], + ikoffset: Field[[ECDim, KDim], int32], + zdiff_gradp: Field[[ECDim, KDim], float], + theta_v_ic: Field[[CellDim, KDim], float], + inv_ddqz_z_full: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + grav_o_cpd: float, + z_hydro_corr: Field[[EdgeDim, KDim], float], ): - closure( - unstructured_domain( - named_range(EdgeDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_solve_nonhydro_stencil_21, - z_hydro_corr, - [ - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, - ], + _mo_solve_nonhydro_stencil_21( + theta_v, + ikoffset, + zdiff_gradp, + theta_v_ic, + inv_ddqz_z_full, + inv_dual_edge_length, + grav_o_cpd, + out=z_hydro_corr, ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "theta_v": FieldInfo( - field=past.FieldSymbol( - id="theta_v", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "ikoffset": FieldInfo( - field=past.FieldSymbol( - id="ikoffset", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zdiff_gradp": FieldInfo( - field=past.FieldSymbol( - id="zdiff_gradp", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "theta_v_ic": FieldInfo( - field=past.FieldSymbol( - id="theta_v_ic", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "inv_ddqz_z_full": FieldInfo( - field=past.FieldSymbol( - id="inv_ddqz_z_full", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "inv_dual_edge_length": FieldInfo( - field=past.FieldSymbol( - id="inv_dual_edge_length", - type=ts.FieldType( - dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "grav_o_cpd": FieldInfo( - field=past.FieldSymbol( - id="grav_o_cpd", - type=ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_hydro_corr": FieldInfo( - field=past.FieldSymbol( - id="z_hydro_corr", - type=ts.FieldType( - dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=False, - out=True, - ), -} - -# patch the fendef with metainfo for icon4pygen -mo_solve_nonhydro_stencil_21.__dict__["offsets"] = [ - Koff.value, - E2C.value, -] # could be done with a pass... -mo_solve_nonhydro_stencil_21.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py new file mode 100644 index 0000000000..b55e3b2194 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py @@ -0,0 +1,200 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.eve import SourceLocation +from gt4py.next.ffront import program_ast as past +from gt4py.next.iterator.builtins import ( + deref, + list_get, + named_range, + power, + shift, + unstructured_domain, +) +from gt4py.next.iterator.runtime import closure, fendef, fundef +from gt4py.next.type_system import type_specifications as ts + +from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff +from icon4py.pyutils.metadata import FieldInfo + + +@fundef +def step(i, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full): + d_ikoffset = list_get(i, deref(ikoffset)) + + d_theta_v = deref(shift(Koff, d_ikoffset, E2C, i)(theta_v)) + s_theta_v_ic = shift(Koff, d_ikoffset, E2C, i)(theta_v_ic) + d_theta_v_ic = deref(s_theta_v_ic) + d_theta_v_ic_p1 = deref(shift(Koff, 1)(s_theta_v_ic)) + d_inv_ddqz_z_full = deref(shift(Koff, d_ikoffset, E2C, i)(inv_ddqz_z_full)) + d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) + + return ( + d_theta_v + d_zdiff_gradp * (d_theta_v_ic - d_theta_v_ic_p1) * d_inv_ddqz_z_full + ) + + +@fundef +def _mo_solve_nonhydro_stencil_21( + theta_v, + ikoffset, + zdiff_gradp, + theta_v_ic, + inv_ddqz_z_full, + inv_dual_edge_length, + grav_o_cpd, +): + z_theta1 = step(0, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) + z_theta2 = step(1, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) + z_hydro_corr = ( + deref(grav_o_cpd) + * deref(inv_dual_edge_length) + * (z_theta2 - z_theta1) + * 4.0 + / power((z_theta1 + z_theta2), 2.0) + ) + return z_hydro_corr + + +@fendef +def mo_solve_nonhydro_stencil_21( + theta_v, + ikoffset, + zdiff_gradp, + theta_v_ic, + inv_ddqz_z_full, + inv_dual_edge_length, + grav_o_cpd, + z_hydro_corr, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, +): + closure( + unstructured_domain( + named_range(EdgeDim, horizontal_start, horizontal_end), + named_range(KDim, vertical_start, vertical_end), + ), + _mo_solve_nonhydro_stencil_21, + z_hydro_corr, + [ + theta_v, + ikoffset, + zdiff_gradp, + theta_v_ic, + inv_ddqz_z_full, + inv_dual_edge_length, + grav_o_cpd, + ], + ) + + +_dummy_loc = SourceLocation(1, 1, "") +_metadata = { + "theta_v": FieldInfo( + field=past.FieldSymbol( + id="theta_v", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "ikoffset": FieldInfo( + field=past.FieldSymbol( + id="ikoffset", + type=ts.FieldType( + dims=[EdgeDim, E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "zdiff_gradp": FieldInfo( + field=past.FieldSymbol( + id="zdiff_gradp", + type=ts.FieldType( + dims=[EdgeDim, E2CDim, KDim], + dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "theta_v_ic": FieldInfo( + field=past.FieldSymbol( + id="theta_v_ic", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "inv_ddqz_z_full": FieldInfo( + field=past.FieldSymbol( + id="inv_ddqz_z_full", + type=ts.FieldType( + dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "inv_dual_edge_length": FieldInfo( + field=past.FieldSymbol( + id="inv_dual_edge_length", + type=ts.FieldType( + dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "grav_o_cpd": FieldInfo( + field=past.FieldSymbol( + id="grav_o_cpd", + type=ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)), + location=_dummy_loc, + ), + inp=True, + out=False, + ), + "z_hydro_corr": FieldInfo( + field=past.FieldSymbol( + id="z_hydro_corr", + type=ts.FieldType( + dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ), + location=_dummy_loc, + ), + inp=False, + out=True, + ), +} + +# patch the fendef with metainfo for icon4pygen +mo_solve_nonhydro_stencil_21.__dict__["offsets"] = [ + Koff.value, + E2C.value, +] # could be done with a pass... +mo_solve_nonhydro_stencil_21.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index bf10a8ea32..496fd4ee87 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -22,6 +22,7 @@ mo_solve_nonhydro_stencil_20_itir, ) +from gt4py.next.ffront.fbuiltins import int32 from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field @@ -76,7 +77,7 @@ def test_mo_solve_nonhydro_stencil_20(): inv_dual_edge_length = random_field(mesh, EdgeDim) z_exner_ex_pr = random_field(mesh, CellDim, KDim) zdiff_gradp = random_field(mesh, EdgeDim, E2CDim, KDim) - ikoffset = zero_field(mesh, EdgeDim, E2CDim, KDim, dtype=int) + ikoffset = zero_field(mesh, EdgeDim, E2CDim, KDim, dtype=int32) rng = np.random.default_rng() for k in range(mesh.k_level): diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py index d98f553268..e8d8c5e2bd 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py @@ -12,14 +12,16 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from gt4py.next.iterator.embedded import constant_field from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_21 import ( mo_solve_nonhydro_stencil_21, ) -from icon4py.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field def mo_solve_nonhydro_stencil_21_numpy( @@ -88,7 +90,7 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): def test_mo_solve_nonhydro_stencil_21(): mesh = SimpleMesh() - ikoffset = zero_field(mesh, EdgeDim, E2CDim, KDim, dtype=int) + ikoffset = zero_field(mesh, EdgeDim, E2CDim, KDim, dtype=int32) rng = np.random.default_rng() for k in range(mesh.k_level): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -105,6 +107,9 @@ def test_mo_solve_nonhydro_stencil_21(): inv_dual_edge_length = random_field(mesh, EdgeDim) grav_o_cpd = 10.0 + zdiff_gradp_new = flatten_first_two_dims(ECDim, KDim, field=zdiff_gradp) + ikoffset_new = flatten_first_two_dims(ECDim, KDim, field=ikoffset) + z_hydro_corr = zero_field(mesh, EdgeDim, KDim) z_hydro_corr_ref = mo_solve_nonhydro_stencil_21_numpy( @@ -125,12 +130,12 @@ def test_mo_solve_nonhydro_stencil_21(): mo_solve_nonhydro_stencil_21( theta_v, - ikoffset, - zdiff_gradp, + ikoffset_new, + zdiff_gradp_new, theta_v_ic, inv_ddqz_z_full, inv_dual_edge_length, - constant_field(grav_o_cpd), + grav_o_cpd, z_hydro_corr, hstart, hend, @@ -138,6 +143,7 @@ def test_mo_solve_nonhydro_stencil_21(): kend, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), "Koff": KDim, }, ) diff --git a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py index 6d9a2d5faa..24628e395e 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py @@ -50,13 +50,13 @@ def render_sid(self) -> str: raise BindingsRenderingException("can not render sid of a scalar") # We want to compute the rank without the sparse dimension, i.e. if a field is horizontal, vertical or both. - # This only works since compound fields are not sparse. - dense_rank = self.entity.rank() - int(self.entity.is_sparse()) - values_str = ( - "1" - if dense_rank == 1 or self.entity.is_compound() - else f"1, mesh_.{self.render_stride_type()}" - ) + dense_rank = self.entity.rank() - int(self.entity.is_sparse() or self.entity.is_compound()) + if dense_rank == 1: + values_str = "1" + elif self.entity.is_compound(): + values_str = f"1, {self.entity.get_num_neighbors()} * mesh_.{self.render_stride_type()}" + else: + values_str = f"1, mesh_.{self.render_stride_type()}" return f"gridtools::hymap::keys<{self.render_dim_tags()}>::make_values({values_str})" def render_ranked_dim_string(self) -> str: @@ -88,11 +88,11 @@ def render_stride_type(self) -> str: _strides = {"E": "EdgeStride", "C": "CellStride", "V": "VertexStride"} if self.entity.is_dense(): return _strides[str(self.entity.location)] - elif self.entity.is_sparse(): + elif self.entity.is_sparse() or self.entity.is_compound(): return _strides[str(self.entity.location[0])] # type: ignore else: raise BindingsRenderingException( - "stride type called on compound location or scalar" + "stride type called on scalar" ) def render_ctype(self, binding_type: str) -> str: diff --git a/pyutils/src/icon4py/icon4pygen/bindings/entities.py b/pyutils/src/icon4py/icon4pygen/bindings/entities.py index 16ae8f6b03..99f6fa5e1d 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/entities.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/entities.py @@ -139,9 +139,9 @@ def rank(self) -> int: return rank def get_num_neighbors(self) -> int: - if not self.is_sparse(): + if not (self.is_sparse() or self.is_compound()): raise BindingsTypeConsistencyException( - "num nbh only defined for sparse fields" + "num nbh only defined for sparse or compound fields" ) return calc_num_neighbors(self.location.to_dim_list(), self.includes_center) # type: ignore diff --git a/pyutils/src/icon4py/icon4pygen/bindings/locations.py b/pyutils/src/icon4py/icon4pygen/bindings/locations.py index 303d776552..e55e305edb 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/locations.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/locations.py @@ -59,6 +59,15 @@ def __init__(self, compound: list[BasicLocation]) -> None: f"chain {compound} contains two of the same elements in succession" ) + def __iter__(self) -> Iterator[BasicLocation]: + return iter(self.compound) + + def __getitem__(self, item: int) -> BasicLocation: + return self.compound[item] + + def to_dim_list(self) -> list[Dimension]: + map_to_dim = {Cell: CellDim, Edge: EdgeDim, Vertex: VertexDim} + return [map_to_dim[c.__class__] for c in self.compound] def is_valid(nbh_list: list[BasicLocation]) -> bool: for i in range(0, len(nbh_list) - 1): # This doesn't look very pythonic From 936d6cca74f1f325f405d10ad879edeac5d66fc7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 2 May 2023 01:42:02 +0200 Subject: [PATCH 07/14] cleanup after rebase on main --- .../mo_nh_diffusion_stencil_15_itir.py | 198 ----------------- .../atm_dyn_iconam/mo_solve_nonhydro_20.py | 0 .../mo_solve_nonhydro_stencil_20_itir.py | 186 ---------------- .../mo_solve_nonhydro_stencil_21_itir.py | 200 ------------------ .../test_mo_solve_nonhydro_stencil_20.py | 29 --- common/src/icon4py/common/dimension.py | 2 + 6 files changed, 2 insertions(+), 613 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py deleted file mode 100644 index db4a421745..0000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15_itir.py +++ /dev/null @@ -1,198 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - if_, - list_get, - named_range, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts - -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo - - -@fundef -def step(i, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset): - d_vcoef = list_get(i, deref(vcoef)) - s_theta_v = shift(C2E2C, i, Koff, list_get(i, deref(zd_vertoffset)))(theta_v) - return list_get(i, deref(geofac_n2s_nbh)) * ( - d_vcoef * deref(s_theta_v) + (1.0 - d_vcoef) * deref(shift(Koff, 1)(s_theta_v)) - ) - - -@fundef -def _mo_nh_diffusion_stencil_15_itir( - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, -): - summed = ( - step(0, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - + step(1, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - + step(2, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset) - ) - update = deref(z_temp) + deref(zd_diffcoef) * ( - deref(theta_v) * deref(geofac_n2s_c) + summed - ) - - return if_(deref(mask), update, deref(z_temp)) - - -@fendef -def mo_nh_diffusion_stencil_15_itir( - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, -): - closure( - unstructured_domain( - named_range(CellDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_nh_diffusion_stencil_15_itir, - z_temp, - [ - mask, - zd_vertoffset, - zd_diffcoef, - geofac_n2s_c, - geofac_n2s_nbh, - vcoef, - theta_v, - z_temp, - ], - ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "mask": FieldInfo( - field=past.FieldSymbol( - id="mask", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zd_vertoffset": FieldInfo( - field=past.FieldSymbol( - id="zd_vertoffset", - type=ts.FieldType( - dims=[CellDim, C2E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zd_diffcoef": FieldInfo( - field=past.FieldSymbol( - id="zd_diffcoef", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "geofac_n2s_c": FieldInfo( - field=past.FieldSymbol( - id="geofac_n2s_c", - type=ts.FieldType( - dims=[CellDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "geofac_n2s_nbh": FieldInfo( - field=past.FieldSymbol( - id="geofac_n2s_nbh", - type=ts.FieldType( - dims=[CellDim, C2E2CDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "vcoef": FieldInfo( - field=past.FieldSymbol( - id="vcoef", - type=ts.FieldType( - dims=[CellDim, C2E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "theta_v": FieldInfo( - field=past.FieldSymbol( - id="theta_v", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_temp": FieldInfo( - field=past.FieldSymbol( - id="z_temp", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=True, - ), -} - -# patch the fendef with metainfo for icon4pygen -mo_nh_diffusion_stencil_15_itir.__dict__["offsets"] = [ - Koff.value, - C2E2C.value, -] # could be done with a pass... -mo_nh_diffusion_stencil_15_itir.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_20.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py deleted file mode 100644 index 2bcd63c5c1..0000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20_itir.py +++ /dev/null @@ -1,186 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - list_get, - named_range, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts - -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo - - -@fundef -def step( - i, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, -): - d_ikoffset = list_get(i, deref(ikoffset)) - d_z_exner_exp_pr = deref(shift(Koff, d_ikoffset, E2C, i)(z_exner_ex_pr)) - d_z_dexner_dz_c_1 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_1)) - d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) - d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) - - return d_z_exner_exp_pr + d_zdiff_gradp * ( - d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 - ) - - -@fundef -def _mo_solve_nonhydro_stencil_20_itir( - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, -): - return deref(inv_dual_edge_length) * ( - step(1, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2) - - step( - 0, z_exner_ex_pr, zdiff_gradp, ikoffset, z_dexner_dz_c_1, z_dexner_dz_c_2 - ) - ) - - -@fendef -def mo_solve_nonhydro_stencil_20_itir( - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, - z_gradh_exner, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, -): - closure( - unstructured_domain( - named_range(EdgeDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_solve_nonhydro_stencil_20_itir, - z_gradh_exner, - [ - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, - ], - ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "inv_dual_edge_length": FieldInfo( - field=past.FieldSymbol( - id="inv_dual_edge_length", - type=ts.FieldType( - dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_exner_ex_pr": FieldInfo( - field=past.FieldSymbol( - id="z_exner_ex_pr", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zdiff_gradp": FieldInfo( - field=past.FieldSymbol( - id="zdiff_gradp", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "ikoffset": FieldInfo( - field=past.FieldSymbol( - id="ikoffset", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_dexner_dz_c_1": FieldInfo( - field=past.FieldSymbol( - id="z_dexner_dz_c_1", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_dexner_dz_c_2": FieldInfo( - field=past.FieldSymbol( - id="z_dexner_dz_c_2", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_gradh_exner": FieldInfo( - field=past.FieldSymbol( - id="z_gradh_exner", - type=ts.FieldType( - dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=False, - out=True, - ), -} -# patch the fendef with metainfo for icon4pygen -mo_solve_nonhydro_stencil_20_itir.__dict__["offsets"] = [ - Koff.value, - E2C.value, -] # could be done with a pass... -mo_solve_nonhydro_stencil_20_itir.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py deleted file mode 100644 index b55e3b2194..0000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21_itir.py +++ /dev/null @@ -1,200 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from gt4py.eve import SourceLocation -from gt4py.next.ffront import program_ast as past -from gt4py.next.iterator.builtins import ( - deref, - list_get, - named_range, - power, - shift, - unstructured_domain, -) -from gt4py.next.iterator.runtime import closure, fendef, fundef -from gt4py.next.type_system import type_specifications as ts - -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo - - -@fundef -def step(i, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full): - d_ikoffset = list_get(i, deref(ikoffset)) - - d_theta_v = deref(shift(Koff, d_ikoffset, E2C, i)(theta_v)) - s_theta_v_ic = shift(Koff, d_ikoffset, E2C, i)(theta_v_ic) - d_theta_v_ic = deref(s_theta_v_ic) - d_theta_v_ic_p1 = deref(shift(Koff, 1)(s_theta_v_ic)) - d_inv_ddqz_z_full = deref(shift(Koff, d_ikoffset, E2C, i)(inv_ddqz_z_full)) - d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) - - return ( - d_theta_v + d_zdiff_gradp * (d_theta_v_ic - d_theta_v_ic_p1) * d_inv_ddqz_z_full - ) - - -@fundef -def _mo_solve_nonhydro_stencil_21( - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, -): - z_theta1 = step(0, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) - z_theta2 = step(1, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full) - z_hydro_corr = ( - deref(grav_o_cpd) - * deref(inv_dual_edge_length) - * (z_theta2 - z_theta1) - * 4.0 - / power((z_theta1 + z_theta2), 2.0) - ) - return z_hydro_corr - - -@fendef -def mo_solve_nonhydro_stencil_21( - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, - z_hydro_corr, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, -): - closure( - unstructured_domain( - named_range(EdgeDim, horizontal_start, horizontal_end), - named_range(KDim, vertical_start, vertical_end), - ), - _mo_solve_nonhydro_stencil_21, - z_hydro_corr, - [ - theta_v, - ikoffset, - zdiff_gradp, - theta_v_ic, - inv_ddqz_z_full, - inv_dual_edge_length, - grav_o_cpd, - ], - ) - - -_dummy_loc = SourceLocation(1, 1, "") -_metadata = { - "theta_v": FieldInfo( - field=past.FieldSymbol( - id="theta_v", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "ikoffset": FieldInfo( - field=past.FieldSymbol( - id="ikoffset", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.INT32), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "zdiff_gradp": FieldInfo( - field=past.FieldSymbol( - id="zdiff_gradp", - type=ts.FieldType( - dims=[EdgeDim, E2CDim, KDim], - dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64), - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "theta_v_ic": FieldInfo( - field=past.FieldSymbol( - id="theta_v_ic", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "inv_ddqz_z_full": FieldInfo( - field=past.FieldSymbol( - id="inv_ddqz_z_full", - type=ts.FieldType( - dims=[CellDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "inv_dual_edge_length": FieldInfo( - field=past.FieldSymbol( - id="inv_dual_edge_length", - type=ts.FieldType( - dims=[EdgeDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "grav_o_cpd": FieldInfo( - field=past.FieldSymbol( - id="grav_o_cpd", - type=ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)), - location=_dummy_loc, - ), - inp=True, - out=False, - ), - "z_hydro_corr": FieldInfo( - field=past.FieldSymbol( - id="z_hydro_corr", - type=ts.FieldType( - dims=[EdgeDim, KDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) - ), - location=_dummy_loc, - ), - inp=False, - out=True, - ), -} - -# patch the fendef with metainfo for icon4pygen -mo_solve_nonhydro_stencil_21.__dict__["offsets"] = [ - Koff.value, - E2C.value, -] # could be done with a pass... -mo_solve_nonhydro_stencil_21.__dict__["metadata"] = _metadata diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 496fd4ee87..9314df2975 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -18,10 +18,6 @@ mo_solve_nonhydro_stencil_20, ) -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20_itir import ( - mo_solve_nonhydro_stencil_20_itir, -) - from gt4py.next.ffront.fbuiltins import int32 from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -93,9 +89,7 @@ def test_mo_solve_nonhydro_stencil_20(): z_dexner_dz_c_1 = random_field(mesh, CellDim, KDim) z_dexner_dz_c_2 = random_field(mesh, CellDim, KDim) - z_gradh_exner = zero_field(mesh, EdgeDim, KDim) - z_gradh_exner_itir = zero_field(mesh, EdgeDim, KDim) z_gradh_exner_ref = mo_solve_nonhydro_stencil_20_numpy( mesh.e2c, @@ -112,24 +106,6 @@ def test_mo_solve_nonhydro_stencil_20(): kstart = 0 kend = mesh.k_level - mo_solve_nonhydro_stencil_20_itir( - inv_dual_edge_length, - z_exner_ex_pr, - zdiff_gradp, - ikoffset, - z_dexner_dz_c_1, - z_dexner_dz_c_2, - z_gradh_exner_itir, - hstart, - hend, - kstart, - kend, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - "Koff": KDim, - }, - ) - mo_solve_nonhydro_stencil_20( inv_dual_edge_length, z_exner_ex_pr, @@ -149,9 +125,4 @@ def test_mo_solve_nonhydro_stencil_20(): }, ) - #print(z_gradh_exner_ref) - #print(z_gradh_exner_itir ) - #print(np.asarray(z_gradh_exner)) - - assert np.allclose(z_gradh_exner_itir, z_gradh_exner_ref) assert np.allclose(z_gradh_exner, z_gradh_exner_ref) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 9193bb8b9d..a0be9cc727 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -21,6 +21,7 @@ CellDim = Dimension("Cell") VertexDim = Dimension("Vertex") CEDim = Dimension("CE") +CECDim = Dimension("CEC") ECDim = Dimension("EC") ECVDim = Dimension("ECV") E2CDim = Dimension("E2C", DimensionKind.LOCAL) @@ -39,6 +40,7 @@ V2E = FieldOffset("V2E", source=EdgeDim, target=(VertexDim, V2EDim)) E2V = FieldOffset("E2V", source=VertexDim, target=(EdgeDim, E2VDim)) C2CE = FieldOffset("C2CE", source=CEDim, target=(CellDim, C2EDim)) +C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim)) E2EC = FieldOffset("E2EC", source=ECDim, target=(EdgeDim, E2CDim)) E2ECV = FieldOffset("E2ECV", source=ECVDim, target=(EdgeDim, E2C2VDim)) E2C2V = FieldOffset("E2C2V", source=VertexDim, target=(EdgeDim, E2C2VDim)) From 9cf3b3f924fa4adada4d1b354f00c6a549339ff1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 2 May 2023 01:51:40 +0200 Subject: [PATCH 08/14] Cleanup --- .../tests/test_mo_nh_diffusion_stencil_15.py | 6 ++-- .../test_mo_solve_nonhydro_stencil_20.py | 1 - .../test_mo_solve_nonhydro_stencil_21.py | 1 - testutils/src/icon4py/testutils/utils.py | 30 +------------------ 4 files changed, 3 insertions(+), 35 deletions(-) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 1fc2d65383..96d00721bb 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -14,6 +14,7 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider +from gt4py.next.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( mo_nh_diffusion_stencil_15, @@ -22,9 +23,6 @@ from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import flatten_first_two_dims, random_field, random_mask, zero_field -from gt4py.next.iterator import embedded as it_embedded - - def mo_nh_diffusion_stencil_15_numpy( c2e2c: np.array, mask: np.array, @@ -69,7 +67,7 @@ def test_mo_nh_diffusion_stencil_15(): mask = random_mask(mesh, CellDim, KDim) - zd_vertoffset = zero_field(mesh, CellDim, C2E2CDim, KDim, dtype=int) + zd_vertoffset = zero_field(mesh, CellDim, C2E2CDim, KDim, dtype=int32) rng = np.random.default_rng() for k in range(mesh.k_level): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 9314df2975..58c7ecc417 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -23,7 +23,6 @@ from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field - def mo_solve_nonhydro_stencil_20_numpy( e2c: np.array, inv_dual_edge_length: np.array, diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py index e8d8c5e2bd..e772ba0f6c 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py @@ -13,7 +13,6 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from gt4py.next.iterator.embedded import constant_field from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_21 import ( mo_solve_nonhydro_stencil_21, diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index eb4e497d3f..6b3cc97974 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -87,51 +87,23 @@ def as_1D_sparse_field( field: it_embedded.MutableLocatedField, dim: gt_common.Dimension ) -> it_embedded.MutableLocatedField: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" - #print("OldArray") - #print(np.asarray(field)) old_shape = np.asarray(field).shape - #print(old_shape) assert len(old_shape) == 2 new_shape = (old_shape[0] * old_shape[1],) - #print("NEWArray") - #print(np.asarray(field).reshape(new_shape)) - #print(new_shape) return it_embedded.np_as_located_field(dim)(np.asarray(field).reshape(new_shape)) def flatten_first_two_dims( *dims: gt_common.Dimension, field: it_embedded.MutableLocatedField ) -> it_embedded.MutableLocatedField: - """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" + """Convert a n-D sparse field to a (n-1)-D flattened (Felix-style) sparse field.""" old_shape = np.asarray(field).shape - #print("Old Shape") - #print(old_shape) assert len(old_shape) >= 2 flattened_size = old_shape[0] * old_shape[1] - #print("Flattened Size") - #print(flattened_size) flattened_shape = (flattened_size, ) - #print("Flattened Shape") - #print(flattened_shape) residual_shape = old_shape[2:] - #print("Residual Shape") - #print(residual_shape) new_shape = flattened_shape + residual_shape - #print("New Shape") - #print(new_shape) - - #print("Field") - #print(field) - array = np.asarray(field) - #print("ARRAY") - #print(array) - #print(array.shape) newarray = np.asarray(field).reshape(new_shape) - #print("NEWARRAY") - #print(newarray) - #print(newarray.shape) - #print(dims) return it_embedded.np_as_located_field(*dims)(newarray) - def get_stencil_module_path(stencil_module: str, stencil_name: str) -> str: return f"icon4py.{stencil_module}.{stencil_name}:{stencil_name}" From 1af19f5bdff448de6d839a192bd12681031d3aba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 2 May 2023 01:53:47 +0200 Subject: [PATCH 09/14] Pre-commit cleanup --- .../mo_nh_diffusion_stencil_15.py | 27 +++++++++--- .../mo_solve_nonhydro_stencil_20.py | 42 +++++++++++++------ .../mo_solve_nonhydro_stencil_21.py | 29 ++++++++++--- .../tests/test_mo_nh_diffusion_stencil_15.py | 11 +++-- .../test_mo_solve_nonhydro_stencil_20.py | 10 +++-- .../test_mo_solve_nonhydro_stencil_21.py | 8 +++- .../bindings/codegen/render/field.py | 8 ++-- .../icon4py/icon4pygen/bindings/locations.py | 1 + testutils/src/icon4py/testutils/utils.py | 4 +- 9 files changed, 104 insertions(+), 36 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 71279d98b8..d3b1ad2c6f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -12,9 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, where, int32 from gt4py.next.ffront.experimental import as_offset -from icon4py.common.dimension import C2E2C, CellDim, CECDim, C2CEC, KDim, Koff +from gt4py.next.ffront.fbuiltins import Field, int32, where + +from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim, Koff + @field_operator def _mo_nh_diffusion_stencil_15( @@ -36,14 +38,29 @@ def _mo_nh_diffusion_stencil_15( theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + int32(1))) theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + int32(1))) - sum = geofac_n2s_nbh(C2CEC[0]) * (vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0])) \ - + geofac_n2s_nbh(C2CEC[1]) * (vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1])) \ - + geofac_n2s_nbh(C2CEC[2]) * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) + sum = ( + geofac_n2s_nbh(C2CEC[0]) + * ( + vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0]) + ) + + geofac_n2s_nbh(C2CEC[1]) + * ( + vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1]) + ) + + geofac_n2s_nbh(C2CEC[2]) + * ( + vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2]) + ) + ) z_temp = where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum), z_temp) return z_temp + @program def mo_nh_diffusion_stencil_15( mask: Field[[CellDim, KDim], bool], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index fb3abd55b4..954da52a47 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -12,10 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum, int32 from gt4py.next.ffront.experimental import as_offset +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2EC, + CellDim, + ECDim, + EdgeDim, + KDim, + Koff, +) -from icon4py.common.dimension import E2C, CellDim, E2EC, ECDim, EdgeDim, KDim, Koff @field_operator def _mo_solve_nonhydro_stencil_20( @@ -36,17 +45,24 @@ def _mo_solve_nonhydro_stencil_20( z_dexner_dz_c2_0 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[0]))) z_dexner_dz_c2_1 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[1]))) - z_gradh_exner = inv_dual_edge_length * (( - z_exner_ex_pr_1(E2C[1]) + - zdiff_gradp(E2EC[1]) * ( - z_dexner_dz_c1_1(E2C[1]) + - zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1]) - )) - ( - z_exner_ex_pr_0(E2C[0]) + - zdiff_gradp(E2EC[0]) * ( - z_dexner_dz_c1_0(E2C[0]) + - zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) - ))) + z_gradh_exner = inv_dual_edge_length * ( + ( + z_exner_ex_pr_1(E2C[1]) + + zdiff_gradp(E2EC[1]) + * ( + z_dexner_dz_c1_1(E2C[1]) + + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1]) + ) + ) + - ( + z_exner_ex_pr_0(E2C[0]) + + zdiff_gradp(E2EC[0]) + * ( + z_dexner_dz_c1_0(E2C[0]) + + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) + ) + ) + ) return z_gradh_exner diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py index f6409ed5c3..2c478358ee 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py @@ -12,10 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum, int32 from gt4py.next.ffront.experimental import as_offset +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2EC, + CellDim, + ECDim, + EdgeDim, + KDim, + Koff, +) -from icon4py.common.dimension import E2C, CellDim, E2EC, ECDim, EdgeDim, KDim, Koff @field_operator def _mo_solve_nonhydro_stencil_21( @@ -40,10 +49,20 @@ def _mo_solve_nonhydro_stencil_21( inv_ddqz_z_full_0 = inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[0]))) inv_ddqz_z_full_1 = inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[1]))) - z_theta_0 = theta_v_0(E2C[0]) + zdiff_gradp(E2EC[0]) * (theta_v_ic_0(E2C[0]) - theta_v_ic_p1_0(E2C[0])) * inv_ddqz_z_full_0(E2C[0]) - z_theta_1 = theta_v_1(E2C[1]) + zdiff_gradp(E2EC[1]) * (theta_v_ic_1(E2C[1]) - theta_v_ic_p1_1(E2C[1])) * inv_ddqz_z_full_1(E2C[1]) + z_theta_0 = theta_v_0(E2C[0]) + zdiff_gradp(E2EC[0]) * ( + theta_v_ic_0(E2C[0]) - theta_v_ic_p1_0(E2C[0]) + ) * inv_ddqz_z_full_0(E2C[0]) + z_theta_1 = theta_v_1(E2C[1]) + zdiff_gradp(E2EC[1]) * ( + theta_v_ic_1(E2C[1]) - theta_v_ic_p1_1(E2C[1]) + ) * inv_ddqz_z_full_1(E2C[1]) - z_hydro_corr = grav_o_cpd * inv_dual_edge_length * (z_theta_1 - z_theta_0) * float(4.0) / (z_theta_0 + z_theta_1)**2 + z_hydro_corr = ( + grav_o_cpd + * inv_dual_edge_length + * (z_theta_1 - z_theta_0) + * float(4.0) + / (z_theta_0 + z_theta_1) ** 2 + ) return z_hydro_corr diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 96d00721bb..9d4f6661a7 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -12,16 +12,21 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np - -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( mo_nh_diffusion_stencil_15, ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import flatten_first_two_dims, random_field, random_mask, zero_field +from icon4py.testutils.utils import ( + flatten_first_two_dims, + random_field, + random_mask, + zero_field, +) + def mo_nh_diffusion_stencil_15_numpy( c2e2c: np.array, diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 58c7ecc417..61e0513c05 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -12,16 +12,20 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, ) - -from gt4py.next.ffront.fbuiltins import int32 from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field +from icon4py.testutils.utils import ( + flatten_first_two_dims, + random_field, + zero_field, +) + def mo_solve_nonhydro_stencil_20_numpy( e2c: np.array, diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py index e772ba0f6c..12b023ee37 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py @@ -12,15 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_21 import ( mo_solve_nonhydro_stencil_21, ) -from gt4py.next.ffront.fbuiltins import int32 from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import flatten_first_two_dims, random_field, zero_field +from icon4py.testutils.utils import ( + flatten_first_two_dims, + random_field, + zero_field, +) def mo_solve_nonhydro_stencil_21_numpy( diff --git a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py index 24628e395e..6b21fd7476 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py @@ -50,7 +50,9 @@ def render_sid(self) -> str: raise BindingsRenderingException("can not render sid of a scalar") # We want to compute the rank without the sparse dimension, i.e. if a field is horizontal, vertical or both. - dense_rank = self.entity.rank() - int(self.entity.is_sparse() or self.entity.is_compound()) + dense_rank = self.entity.rank() - int( + self.entity.is_sparse() or self.entity.is_compound() + ) if dense_rank == 1: values_str = "1" elif self.entity.is_compound(): @@ -91,9 +93,7 @@ def render_stride_type(self) -> str: elif self.entity.is_sparse() or self.entity.is_compound(): return _strides[str(self.entity.location[0])] # type: ignore else: - raise BindingsRenderingException( - "stride type called on scalar" - ) + raise BindingsRenderingException("stride type called on scalar") def render_ctype(self, binding_type: str) -> str: """Render C datatype for a corresponding binding type.""" diff --git a/pyutils/src/icon4py/icon4pygen/bindings/locations.py b/pyutils/src/icon4py/icon4pygen/bindings/locations.py index e55e305edb..84f6e149a0 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/locations.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/locations.py @@ -69,6 +69,7 @@ def to_dim_list(self) -> list[Dimension]: map_to_dim = {Cell: CellDim, Edge: EdgeDim, Vertex: VertexDim} return [map_to_dim[c.__class__] for c in self.compound] + def is_valid(nbh_list: list[BasicLocation]) -> bool: for i in range(0, len(nbh_list) - 1): # This doesn't look very pythonic if isinstance(type(nbh_list[i]), type(nbh_list[i + 1])): diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index 6b3cc97974..ad5c5128d5 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -92,6 +92,7 @@ def as_1D_sparse_field( new_shape = (old_shape[0] * old_shape[1],) return it_embedded.np_as_located_field(dim)(np.asarray(field).reshape(new_shape)) + def flatten_first_two_dims( *dims: gt_common.Dimension, field: it_embedded.MutableLocatedField ) -> it_embedded.MutableLocatedField: @@ -99,11 +100,12 @@ def flatten_first_two_dims( old_shape = np.asarray(field).shape assert len(old_shape) >= 2 flattened_size = old_shape[0] * old_shape[1] - flattened_shape = (flattened_size, ) + flattened_shape = (flattened_size,) residual_shape = old_shape[2:] new_shape = flattened_shape + residual_shape newarray = np.asarray(field).reshape(new_shape) return it_embedded.np_as_located_field(*dims)(newarray) + def get_stencil_module_path(stencil_module: str, stencil_name: str) -> str: return f"icon4py.{stencil_module}.{stencil_name}:{stencil_name}" From 3cd049fc7a4f187912d6e31869b494aa780425dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 5 May 2023 03:10:56 +0200 Subject: [PATCH 10/14] Created abstract baseclass for ChainedLocation and CompoundLocation --- .../icon4py/icon4pygen/bindings/locations.py | 43 +++++++------------ 1 file changed, 15 insertions(+), 28 deletions(-) diff --git a/pyutils/src/icon4py/icon4pygen/bindings/locations.py b/pyutils/src/icon4py/icon4pygen/bindings/locations.py index 84f6e149a0..6ec4afb88e 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/locations.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/locations.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from abc import ABCMeta, abstractmethod from typing import Iterator from gt4py.next.ffront.fbuiltins import Dimension @@ -45,31 +46,6 @@ def __str__(self) -> str: BASIC_LOCATIONS = {location.__name__: location for location in [Cell, Edge, Vertex]} -class CompoundLocation: - compound: list[BasicLocation] - - def __str__(self) -> str: - return "".join([str(loc) for loc in self.compound]) - - def __init__(self, compound: list[BasicLocation]) -> None: - if is_valid(compound): - self.compound = compound - else: - raise Exception( - f"chain {compound} contains two of the same elements in succession" - ) - - def __iter__(self) -> Iterator[BasicLocation]: - return iter(self.compound) - - def __getitem__(self, item: int) -> BasicLocation: - return self.compound[item] - - def to_dim_list(self) -> list[Dimension]: - map_to_dim = {Cell: CellDim, Edge: EdgeDim, Vertex: VertexDim} - return [map_to_dim[c.__class__] for c in self.compound] - - def is_valid(nbh_list: list[BasicLocation]) -> bool: for i in range(0, len(nbh_list) - 1): # This doesn't look very pythonic if isinstance(type(nbh_list[i]), type(nbh_list[i + 1])): @@ -77,12 +53,12 @@ def is_valid(nbh_list: list[BasicLocation]) -> bool: return True -class ChainedLocation: +class MultiLocation(metaclass=ABCMeta): chain: list[BasicLocation] + @abstractmethod def __str__(self) -> str: - return "2".join([str(loc) for loc in self.chain]) - + ... def __init__(self, chain: list[BasicLocation]) -> None: if is_valid(chain): self.chain = chain @@ -100,3 +76,14 @@ def __getitem__(self, item: int) -> BasicLocation: def to_dim_list(self) -> list[Dimension]: map_to_dim = {Cell: CellDim, Edge: EdgeDim, Vertex: VertexDim} return [map_to_dim[c.__class__] for c in self.chain] + +class CompoundLocation(MultiLocation): + + def __str__(self) -> str: + return "".join([str(loc) for loc in self.chain]) + + +class ChainedLocation(MultiLocation): + + def __str__(self) -> str: + return "2".join([str(loc) for loc in self.chain]) From 362b9e870c777ac23383615e4e4ec5a79a185ed0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 8 May 2023 14:00:39 +0200 Subject: [PATCH 11/14] Make pre-commit happy --- .../src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py | 4 ++-- .../icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py | 2 +- .../icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index d3b1ad2c6f..5768729552 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -38,7 +38,7 @@ def _mo_nh_diffusion_stencil_15( theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + int32(1))) theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + int32(1))) - sum = ( + neighbor_sum = ( geofac_n2s_nbh(C2CEC[0]) * ( vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) @@ -56,7 +56,7 @@ def _mo_nh_diffusion_stencil_15( ) ) - z_temp = where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum), z_temp) + z_temp = where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + neighbor_sum), z_temp) return z_temp diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index 954da52a47..643c4c9b3e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -13,7 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.experimental import as_offset -from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import ( E2C, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py index 2c478358ee..5e42edfe24 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py @@ -13,7 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.experimental import as_offset -from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import ( E2C, From 59e62c2d00a715414cd40a9dae237b655fc592ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 9 May 2023 00:24:30 +0200 Subject: [PATCH 12/14] Formatting fix --- .../src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py | 4 +++- pyutils/src/icon4py/icon4pygen/bindings/locations.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 5768729552..226c461aa7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -56,7 +56,9 @@ def _mo_nh_diffusion_stencil_15( ) ) - z_temp = where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + neighbor_sum), z_temp) + z_temp = where( + mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + neighbor_sum), z_temp + ) return z_temp diff --git a/pyutils/src/icon4py/icon4pygen/bindings/locations.py b/pyutils/src/icon4py/icon4pygen/bindings/locations.py index 6ec4afb88e..8222a0c20c 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/locations.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/locations.py @@ -59,6 +59,7 @@ class MultiLocation(metaclass=ABCMeta): @abstractmethod def __str__(self) -> str: ... + def __init__(self, chain: list[BasicLocation]) -> None: if is_valid(chain): self.chain = chain @@ -77,13 +78,12 @@ def to_dim_list(self) -> list[Dimension]: map_to_dim = {Cell: CellDim, Edge: EdgeDim, Vertex: VertexDim} return [map_to_dim[c.__class__] for c in self.chain] -class CompoundLocation(MultiLocation): +class CompoundLocation(MultiLocation): def __str__(self) -> str: return "".join([str(loc) for loc in self.chain]) class ChainedLocation(MultiLocation): - def __str__(self) -> str: return "2".join([str(loc) for loc in self.chain]) From fd3f983314a1ae916bdc1be960dd86bfc51a468c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 10 May 2023 14:24:45 +0200 Subject: [PATCH 13/14] Address Nikki's feedback --- .../src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py | 4 ++-- testutils/src/icon4py/testutils/utils.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 226c461aa7..6a142e5de6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -38,7 +38,7 @@ def _mo_nh_diffusion_stencil_15( theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + int32(1))) theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + int32(1))) - neighbor_sum = ( + sum_over_neighbors = ( geofac_n2s_nbh(C2CEC[0]) * ( vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) @@ -57,7 +57,7 @@ def _mo_nh_diffusion_stencil_15( ) z_temp = where( - mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + neighbor_sum), z_temp + mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over_neighbors), z_temp ) return z_temp diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index ad5c5128d5..9726a0f064 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -101,8 +101,7 @@ def flatten_first_two_dims( assert len(old_shape) >= 2 flattened_size = old_shape[0] * old_shape[1] flattened_shape = (flattened_size,) - residual_shape = old_shape[2:] - new_shape = flattened_shape + residual_shape + new_shape = flattened_shape + old_shape[2:] newarray = np.asarray(field).reshape(new_shape) return it_embedded.np_as_located_field(*dims)(newarray) From b245b90b2771cd6c2accc572be941393664401ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 10 May 2023 14:25:52 +0200 Subject: [PATCH 14/14] formatting --- .../src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 6a142e5de6..9abed41ee1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -57,7 +57,9 @@ def _mo_nh_diffusion_stencil_15( ) z_temp = where( - mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over_neighbors), z_temp + mask, + z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over_neighbors), + z_temp, ) return z_temp