From ff90c8461fcb2dc4bc9ad33b7648df93876b2866 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 28 Apr 2023 10:50:45 +0200 Subject: [PATCH 1/5] Workaround a change in GT4Py (#197) Explanation: - removing `@fundef` will change how the tracing of the iterator IR program works: it will execute the function and basically inline on tracing. - after it's inlined in iterator IR, the `i` argument of `step` will be a `itir.OffsetLiteral` in the shift (instead of a `itir.Literal` - this will then correctly translate to `integral_constant`s (while the `Literal` would translate to `int`), which is not legal in the unstructured shift. All this is triggered by cleaning up the `integral_constant` handling in https://github.com/GridTools/gt4py/pull/1204. --- .../src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py | 1 - 1 file changed, 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 345bd03db1..a0d3fca5b1 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 @@ -28,7 +28,6 @@ from icon4py.icon4pygen.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) From 37505eaa8cce149df63f8a235d8b613d272ce5cf Mon Sep 17 00:00:00 2001 From: juckerj <39263956+jonasjucker@users.noreply.github.com> Date: Tue, 2 May 2023 10:01:12 +0200 Subject: [PATCH 2/5] [Jenkins] Daily plan with Spack (#191) --- jenkins/spack-daily | 33 +++++++++++++++++++ .../src/icon4py/icon4pygen/bindings/utils.py | 13 +++++--- 2 files changed, 41 insertions(+), 5 deletions(-) create mode 100644 jenkins/spack-daily diff --git a/jenkins/spack-daily b/jenkins/spack-daily new file mode 100644 index 0000000000..852724ad32 --- /dev/null +++ b/jenkins/spack-daily @@ -0,0 +1,33 @@ +pipeline { + agent none + stages { + stage('Tests') { + matrix { + agent { label "${NODENAME}" } + axes { + axis { + name 'NODENAME' + values 'daint', 'balfrin' + } + } + post { + always { + echo 'Cleaning up workspace' + deleteDir() + } + } + stages { + stage('Install and Test') { + steps { + sh """ + git clone --depth 1 --recurse-submodules --shallow-submodules https://github.com/C2SM/spack-c2sm.git + . ./spack-c2sm/setup-env.sh + spack install -v --test=root py-icon4py@main%gcc ^py-gt4py@main%gcc + """ + } + } + } + } + } + } +} diff --git a/pyutils/src/icon4py/icon4pygen/bindings/utils.py b/pyutils/src/icon4py/icon4pygen/bindings/utils.py index ca0c47c4ab..d33c56a987 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/utils.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/utils.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import os import subprocess import sys from pathlib import Path @@ -42,13 +43,15 @@ def format_fortran_code(source: str) -> str: The path to fprettify needs to be set explicitly for the non-Spack build process as Liskov does not activate a virtual environment. - However, the PYTHON_PATH does not contain fprettify in a Spack build, hence the need for a special condition + + Variable SPACK_ROOT is always set in a spack build, used to assume that fprettify + is in PATH """ - bin_path = Path(PYTHON_PATH).parent - if "spack" not in str(bin_path): - fprettify_path = bin_path / "fprettify" - else: + if os.getenv("SPACK_ROOT") is not None: fprettify_path = "fprettify" + else: + bin_path = Path(PYTHON_PATH).parent + fprettify_path = bin_path / "fprettify" args = [str(fprettify_path)] p1 = subprocess.Popen(args, stdout=subprocess.PIPE, stdin=subprocess.PIPE) return p1.communicate(source.encode("UTF-8"))[0].decode("UTF-8").rstrip() From 51236912ef80a7c89bb758f663b8d5187a46f52d Mon Sep 17 00:00:00 2001 From: juckerj <39263956+jonasjucker@users.noreply.github.com> Date: Thu, 4 May 2023 16:35:12 +0200 Subject: [PATCH 3/5] use shutil.which to find fprettify (#203) --- pyutils/src/icon4py/icon4pygen/bindings/utils.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/pyutils/src/icon4py/icon4pygen/bindings/utils.py b/pyutils/src/icon4py/icon4pygen/bindings/utils.py index d33c56a987..83d2770d1c 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/utils.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/utils.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import os +import shutil import subprocess import sys from pathlib import Path @@ -41,15 +41,12 @@ def calc_num_neighbors(dim_list: list[Dimension], includes_center: bool) -> int: def format_fortran_code(source: str) -> str: """Format fortran code using fprettify. - The path to fprettify needs to be set explicitly for the - non-Spack build process as Liskov does not activate a virtual environment. - - Variable SPACK_ROOT is always set in a spack build, used to assume that fprettify - is in PATH + Try to find fprettify in PATH -> found by which + otherwise look in PYTHONPATH """ - if os.getenv("SPACK_ROOT") is not None: - fprettify_path = "fprettify" - else: + fprettify_path = shutil.which("fprettify") + + if fprettify_path is None: bin_path = Path(PYTHON_PATH).parent fprettify_path = bin_path / "fprettify" args = [str(fprettify_path)] From eea3bf8f597c6b40fd2ee03b91b0c49da9e80488 Mon Sep 17 00:00:00 2001 From: muellch <60387010+muellch@users.noreply.github.com> Date: Mon, 15 May 2023 10:08:59 +0200 Subject: [PATCH 4/5] Index stencils field view implementation (#199) Write field view representation for offset stencils * Adapt bindings generator to be able to handle compound fields who also have a sparse dimension (example `Field[[CEDim, KDim], float]`. * Introduced a `flatten_first_two_dims` helper method to flatten the first two dimensions of a field, which need to be the two dimensions representing a dense and a sparse dimension (example `Field[[EdgeDim, E2CDim, ...], float]`), into a compound dimension (example `Field[[ECDim, ...], float]`). * Refactor sparse and compound location type, using a common base class there. --- .../mo_nh_diffusion_stencil_15.py | 232 +++++------------- .../mo_solve_nonhydro_stencil_20.py | 220 +++++------------ .../mo_solve_nonhydro_stencil_21.py | 232 +++++------------- .../tests/test_mo_nh_diffusion_stencil_15.py | 24 +- .../test_mo_solve_nonhydro_stencil_20.py | 22 +- .../test_mo_solve_nonhydro_stencil_21.py | 23 +- common/src/icon4py/common/dimension.py | 2 + .../bindings/codegen/render/field.py | 20 +- .../icon4py/icon4pygen/bindings/entities.py | 4 +- .../icon4py/icon4pygen/bindings/locations.py | 31 ++- testutils/src/icon4py/testutils/utils.py | 13 + 11 files changed, 275 insertions(+), 548 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 a0d3fca5b1..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 @@ -11,187 +11,79 @@ # # 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 gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.experimental import as_offset +from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, Koff -from icon4py.icon4pygen.metadata import FieldInfo +from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim, Koff -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)) - ) +@field_operator +def _mo_nh_diffusion_stencil_15( + mask: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim, KDim], int32], + 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]: + 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]))) -@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) + 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_over_neighbors = ( + 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]) + ) ) - update = deref(z_temp) + deref(zd_diffcoef) * ( - deref(theta_v) * deref(geofac_n2s_c) + summed + + z_temp = where( + mask, + z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over_neighbors), + z_temp, ) - return if_(deref(mask), update, deref(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], int32], + 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_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index 06a3829b29..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 @@ -11,176 +11,78 @@ # # 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.ffront.decorator import field_operator, program +from gt4py.next.ffront.experimental import as_offset +from gt4py.next.ffront.fbuiltins import Field, int32 + +from icon4py.common.dimension import ( + E2C, + E2EC, + CellDim, + ECDim, + EdgeDim, + KDim, + Koff, ) -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.icon4pygen.metadata import FieldInfo +@field_operator +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], int32], + z_dexner_dz_c_1: Field[[CellDim, KDim], float], + z_dexner_dz_c_2: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: -@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)) + 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]))) - return d_z_exner_exp_pr + d_zdiff_gradp * ( - d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 - ) + 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]))) -@fundef -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 + 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 + -@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[[ECDim, KDim], float], + 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], ): - 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 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..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 @@ -11,190 +11,80 @@ # # 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.ffront.decorator import field_operator, program +from gt4py.next.ffront.experimental import as_offset +from gt4py.next.ffront.fbuiltins import Field, int32 + +from icon4py.common.dimension import ( + E2C, + E2EC, + CellDim, + ECDim, + EdgeDim, + KDim, + Koff, ) -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.icon4pygen.metadata import FieldInfo +@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]) -@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) + 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 -@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/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 829fb9391c..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,13 +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_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, +) def mo_nh_diffusion_stencil_15_numpy( @@ -65,7 +72,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`) @@ -82,6 +89,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 +112,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 +125,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 040a88ba25..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,13 +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_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( @@ -70,7 +76,8 @@ 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): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -80,9 +87,11 @@ 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) - z_gradh_exner = zero_field(mesh, EdgeDim, KDim) z_gradh_exner_ref = mo_solve_nonhydro_stencil_20_numpy( @@ -103,8 +112,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,6 +123,7 @@ 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, }, ) 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..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,14 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.iterator.embedded import constant_field +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 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_21_numpy( @@ -88,7 +93,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 +110,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 +133,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 +146,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/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)) diff --git a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py index 6d9a2d5faa..6b21fd7476 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py @@ -50,13 +50,15 @@ 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,12 +90,10 @@ 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" - ) + 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/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..8222a0c20c 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,21 +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 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])): @@ -67,11 +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): @@ -90,3 +77,13 @@ 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]) diff --git a/testutils/src/icon4py/testutils/utils.py b/testutils/src/icon4py/testutils/utils.py index e4d503cfca..9726a0f064 100644 --- a/testutils/src/icon4py/testutils/utils.py +++ b/testutils/src/icon4py/testutils/utils.py @@ -93,5 +93,18 @@ def as_1D_sparse_field( 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 n-D sparse field to a (n-1)-D flattened (Felix-style) sparse field.""" + old_shape = np.asarray(field).shape + assert len(old_shape) >= 2 + flattened_size = old_shape[0] * old_shape[1] + flattened_shape = (flattened_size,) + new_shape = flattened_shape + old_shape[2:] + 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 be34d204a04e8752410fe891de9d1e9af34bc6fd Mon Sep 17 00:00:00 2001 From: Samuel Date: Tue, 16 May 2023 18:10:23 +0200 Subject: [PATCH 5/5] Introduce __version__ (#205) Declare version inside `liskov` --- liskov/src/icon4py/liskov/__init__.py | 9 ++++ liskov/src/icon4py/liskov/codegen/f90.py | 3 +- .../src/icon4py/liskov/external/exceptions.py | 4 -- .../src/icon4py/liskov/external/metadata.py | 47 ++++--------------- liskov/tests/test_code_metadata.py | 30 ++---------- 5 files changed, 22 insertions(+), 71 deletions(-) diff --git a/liskov/src/icon4py/liskov/__init__.py b/liskov/src/icon4py/liskov/__init__.py index 15dfdb0098..155fd013a9 100644 --- a/liskov/src/icon4py/liskov/__init__.py +++ b/liskov/src/icon4py/liskov/__init__.py @@ -10,3 +10,12 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Final + + +__all__ = [ + "__version__", +] + +__version__: Final = "0.0.3" diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index 0d2b0ed2de..6ecc3e99ff 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -122,8 +122,7 @@ class MetadataStatementGenerator(TemplatedGenerator): ! Generated on: {{ _this_node.metadata.generated_on }} ! Input filepath: {{ _this_node.metadata.cli_params['input_filepath'] }} ! Profiling active: {{ _this_node.metadata.cli_params['profile'] }} - ! Git version tag: {{ _this_node.metadata.tag }} - ! Git commit hash: {{ _this_node.metadata.commit_hash }} + ! Git version tag: {{ _this_node.metadata.version }} !+-+-+-+-+-+-+-+-+-+ +-+-+-+-+ +-+-+-+-+-+-+-+-+-+-+-+ """ ) diff --git a/liskov/src/icon4py/liskov/external/exceptions.py b/liskov/src/icon4py/liskov/external/exceptions.py index 2ae3931925..15f28dadc6 100644 --- a/liskov/src/icon4py/liskov/external/exceptions.py +++ b/liskov/src/icon4py/liskov/external/exceptions.py @@ -16,10 +16,6 @@ class MissingClickContextError(Exception): pass -class MissingGitError(Exception): - pass - - class UnknownStencilError(Exception): pass diff --git a/liskov/src/icon4py/liskov/external/metadata.py b/liskov/src/icon4py/liskov/external/metadata.py index c47fedcc3b..e1b5ff2c47 100644 --- a/liskov/src/icon4py/liskov/external/metadata.py +++ b/liskov/src/icon4py/liskov/external/metadata.py @@ -11,24 +11,20 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import datetime -import subprocess -from pathlib import Path from typing import Any import click -from icon4py.liskov.external.exceptions import ( - MissingClickContextError, - MissingGitError, -) +import icon4py.liskov +from icon4py.liskov.external.exceptions import MissingClickContextError class CodeMetadata: """Class that handles retrieval of icon-liskov runtime metadata.""" - def __init__(self) -> None: - self.generated_on = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - self.parent_dir = Path(__file__).parent + @property + def generated_on(self) -> str: + return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") @property def cli_params(self) -> dict[str, Any]: @@ -41,33 +37,6 @@ def cli_params(self) -> dict[str, Any]: ) @property - def commit_hash(self) -> str: - """Get the latest git commit hash.""" - try: - return ( - subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=self.parent_dir - ) - .decode() - .strip() - ) - except Exception as e: - raise MissingGitError( - f"Git is not available or there is no commit or tag.\n {e}" - ) - - @property - def tag(self) -> str: - """Get the latest git tag.""" - try: - return ( - subprocess.check_output( - ["git", "describe", "--tags", "--abbrev=0"], cwd=self.parent_dir - ) - .decode() - .strip() - ) - except Exception as e: - raise MissingGitError( - f"Git is not available or there is no commit or tag.\n {e}" - ) + def version(self) -> str: + """Get the current version.""" + return icon4py.liskov.__version__ diff --git a/liskov/tests/test_code_metadata.py b/liskov/tests/test_code_metadata.py index 6dfb128007..16c5100fcb 100644 --- a/liskov/tests/test_code_metadata.py +++ b/liskov/tests/test_code_metadata.py @@ -16,7 +16,7 @@ import pytest -from icon4py.liskov.external.exceptions import MissingGitError +import icon4py.liskov from icon4py.liskov.external.metadata import CodeMetadata @@ -34,31 +34,9 @@ def test_generated_on(): assert metadata.generated_on == "2023-03-01 12:00:00" -def test_tag(module_parent): - with mock.patch("subprocess.check_output", side_effect=[b"v1.0.0"]) as mock_git: - metadata = CodeMetadata() - assert metadata.tag == "v1.0.0" - mock_git.assert_called_once_with( - ["git", "describe", "--tags", "--abbrev=0"], cwd=module_parent - ) - - -def test_commit_hash(module_parent): - with mock.patch( - "subprocess.check_output", side_effect=[b"abcdef123456"] - ) as mock_git: - metadata = CodeMetadata() - assert metadata.commit_hash == "abcdef123456" - mock_git.assert_called_once_with( - ["git", "rev-parse", "HEAD"], cwd=module_parent - ) - - -def test_no_git(): - with mock.patch("subprocess.check_output", side_effect=MissingGitError()): - metadata = CodeMetadata() - with pytest.raises(MissingGitError): - metadata.tag +def test_version(module_parent): + metadata = CodeMetadata() + assert metadata.version == icon4py.liskov.__version__ def test_click_context():