From 2e37841d2832d5bc9d23fcd7c99cbdd40314b77f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 9 Sep 2022 17:35:00 +0200 Subject: [PATCH 001/105] Add first 5 tracer advection stencils to icon4py Co-authored-by: Erwan Cossevin Co-authored-by: Stephanie Westerhuis --- .coveragerc | 4 + advection/README.md | 9 +++ advection/requirements-dev.txt | 3 + advection/requirements.txt | 3 + advection/setup.cfg | 50 +++++++++++++ advection/setup.py | 17 +++++ advection/src/icon4py/advection/__init__.py | 13 ++++ .../advection/hflx_limiter_pd_stencil_01.py | 56 ++++++++++++++ .../advection/hflx_limiter_pd_stencil_02.py | 51 +++++++++++++ .../advection/rbf_intp_edge_stencil_01.py | 35 +++++++++ .../advection/step_advection_stencil_04.py | 35 +++++++++ .../tests/test_hflx_limiter_pd_stencil_01.py | 74 +++++++++++++++++++ .../tests/test_hflx_limiter_pd_stencil_02.py | 70 ++++++++++++++++++ .../tests/test_rbf_intp_edge_stencil_01.py | 57 ++++++++++++++ .../tests/test_step_advection_stencil_04.py | 52 +++++++++++++ .../icon4py/atm_dyn_iconam/compute_airmass.py | 37 ++++++++++ atm_dyn_iconam/tests/test_compute_airmass.py | 44 +++++++++++ requirements-dev.txt | 1 + requirements.txt | 1 + 19 files changed, 612 insertions(+) create mode 100644 advection/README.md create mode 100644 advection/requirements-dev.txt create mode 100644 advection/requirements.txt create mode 100644 advection/setup.cfg create mode 100644 advection/setup.py create mode 100644 advection/src/icon4py/advection/__init__.py create mode 100644 advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py create mode 100644 advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py create mode 100644 advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py create mode 100644 advection/src/icon4py/advection/step_advection_stencil_04.py create mode 100644 advection/tests/test_hflx_limiter_pd_stencil_01.py create mode 100644 advection/tests/test_hflx_limiter_pd_stencil_02.py create mode 100644 advection/tests/test_rbf_intp_edge_stencil_01.py create mode 100644 advection/tests/test_step_advection_stencil_04.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py create mode 100644 atm_dyn_iconam/tests/test_compute_airmass.py diff --git a/.coveragerc b/.coveragerc index e577d7bd05..2c097ba2c0 100644 --- a/.coveragerc +++ b/.coveragerc @@ -11,6 +11,10 @@ atm_dyn_iconam_sources = atm_dyn_iconam/src/icon4py/atm_dyn_iconam/ .tox/py*/lib/python3.*/site-packages/icon4py/atm_dyn_iconam/ +advection_sources = + advection/src/icon4py/advection/ + .tox/py*/lib/python3.*/site-packages/icon4py/advection/ + common_sources = common/src/icon4py/common/ .tox/py*/lib/python3.*/site-packages/icon4py/common/ diff --git a/advection/README.md b/advection/README.md new file mode 100644 index 0000000000..9ee5f3fdab --- /dev/null +++ b/advection/README.md @@ -0,0 +1,9 @@ +# icon4py-advection + +## Description + +Contains code ported from ICON `src/advection`. + +## Installation instructions + +Check `README.md` file in the root of the repository. diff --git a/advection/requirements-dev.txt b/advection/requirements-dev.txt new file mode 100644 index 0000000000..2cb7bda5e8 --- /dev/null +++ b/advection/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../base-requirements-dev.txt +-e ../common +-e . diff --git a/advection/requirements.txt b/advection/requirements.txt new file mode 100644 index 0000000000..c8420e5bf8 --- /dev/null +++ b/advection/requirements.txt @@ -0,0 +1,3 @@ +-r ../base-requirements.txt +../common +. diff --git a/advection/setup.cfg b/advection/setup.cfg new file mode 100644 index 0000000000..0d8a4b9166 --- /dev/null +++ b/advection/setup.cfg @@ -0,0 +1,50 @@ +# This file is mainly used to configure package creation with setuptools. +# Documentation: +# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files +# +[metadata] +name = icon4py_advection +description = Icon inspired code in Python and GT4Py +long_description = file: README.md +long_description_content_type = text/markdown +url = https://github.com/C2SM/icon4py +author = ETH Zurich +author_email = gridtools@cscs.ch +license = gpl3 +license_files = LICENSE +platforms = Linux, Mac +classifiers = + Development Status :: 3 - Alpha + Intended Audience :: Science/Research + License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) + Operating System :: POSIX + Programming Language :: Python + Programming Language :: Python :: 3 + Programming Language :: Python :: 3 :: Only + Programming Language :: Python :: 3.10 + Programming Language :: Python :: Implementation :: CPython + Topic :: Scientific/Engineering :: Atmospheric Science + Topic :: Scientific/Engineering :: Mathematics + Topic :: Scientific/Engineering :: Physics +project_urls = + Source Code = https://github.com/C2SM/icon4py + +[options] +packages = find_namespace: +install_requires = + icon4py-common +python_requires = >=3.10 +package_dir = + = src +zip_safe = False + +[options.package_data] +# References: +# https://setuptools.pypa.io/en/latest/userguide/datafiles.html +# https://github.com/abravalheri/experiment-setuptools-package-data +* = *.md, *.rst, *.toml, *.txt, py.typed + +[options.packages.find] +where = src +exclude = + tests diff --git a/advection/setup.py b/advection/setup.py new file mode 100644 index 0000000000..ceec936ea5 --- /dev/null +++ b/advection/setup.py @@ -0,0 +1,17 @@ +# +# 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 setuptools import setup + + +if __name__ == "__main__": + setup() diff --git a/advection/src/icon4py/advection/__init__.py b/advection/src/icon4py/advection/__init__.py new file mode 100644 index 0000000000..a6d2d236c9 --- /dev/null +++ b/advection/src/icon4py/advection/__init__.py @@ -0,0 +1,13 @@ +# 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 + diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py new file mode 100644 index 0000000000..94c13720b0 --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py @@ -0,0 +1,56 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, maximum, minimum, neighbor_sum, broadcast +from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim, C2E + +@field_operator +def _hflx_limiter_pd_stencil_01( + geofac_div: Field[[CellDim, C2EDim], float], + p_cc: Field[[CellDim, KDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_dtime:float, + dbl_eps: float, +) -> Field[[CellDim, KDim], float]: +# p_m: Field[[CellDim, KDim], float] +# p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) +# p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[1])*geofac_div(1)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[2])*geofac_div(2)*p_dtime) + #p_m = neighbor_sum(a_a, axis=C2EDim) + #r_m = minimum(1.0, (p_cc*p_rhodz_now)/(p_m+dbl_eps)) +# r_m = minimum(1.0, (p_cc*p_rhodz_now)/(1.0+dbl_eps)) + + p_m = neighbor_sum(maximum(broadcast(0.0, (CellDim, KDim)), p_mflx_tracer_h(C2E)*geofac_div*p_dtime), axis=C2EDim) + r_m = minimum(broadcast(1.0, (CellDim, KDim)), (p_cc*p_rhodz_now)/(p_m+dbl_eps)) + return r_m + +@program +def hflx_limiter_pd_stencil_01( + geofac_div: Field[[CellDim, C2EDim], float], + p_cc: Field[[CellDim, KDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + p_dtime:float, + dbl_eps: float, +): + _hflx_limiter_pd_stencil_01( + geofac_div, + p_cc, + p_rhodz_now, + p_mflx_tracer_h, + p_dtime, + dbl_eps, + out=r_m, + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py new file mode 100644 index 0000000000..2285800f89 --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py @@ -0,0 +1,51 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where + +from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2C + +@field_operator +def _hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], float], + r_m: Field[[CellDim, KDim], float], + p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], + bound: float, +) -> Field[[EdgeDim, KDim], float]: + p_mflx_tracer_h_out = where ( + refin_ctrl == bound, + p_mflx_tracer_h_in, + where ( + (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in==0.0), + p_mflx_tracer_h_in * r_m(E2C[0]), + p_mflx_tracer_h_in * r_m(E2C[1]), + ), + ) + return p_mflx_tracer_h_out + +@program +def hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], float], + r_m: Field[[CellDim, KDim], float], + p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h_out: Field[[EdgeDim, KDim], float], + bound: float, +): + _hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + out=p_mflx_tracer_h_out, +) diff --git a/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py b/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py new file mode 100644 index 0000000000..76a97b58c9 --- /dev/null +++ b/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py @@ -0,0 +1,35 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim + + +@field_operator +def _rbf_intp_edge_stencil_01( + p_vn_in: Field[[EdgeDim, KDim], float], + # ptr_coeff: Field[[EdgeDim, E2CDim, E2C2EDim], float] + ptr_coeff: Field[[EdgeDim, E2C2EDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim) + return p_vt_out + +@program +def rbf_intp_edge_stencil_01( + p_vn_in: Field[[EdgeDim, KDim], float], + ptr_coeff: Field[[EdgeDim, E2C2EDim], float], + p_vt_out: Field[[EdgeDim, KDim], float], +): + _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out) diff --git a/advection/src/icon4py/advection/step_advection_stencil_04.py b/advection/src/icon4py/advection/step_advection_stencil_04.py new file mode 100644 index 0000000000..f299d11d98 --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_04.py @@ -0,0 +1,35 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import CellDim, KDim + +@field_operator +def _step_advection_stencil_04( + p_tracer_now: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + opt_ddt_tracer_adv=(p_tracer_new-p_tracer_now)/p_dtime + return opt_ddt_tracer_adv + +@program +def step_advection_stencil_04( + p_tracer_now: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + opt_ddt_tracer_adv: Field[[CellDim, KDim], float], + p_dtime: float, +): + _step_advection_stencil_04(p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/advection/tests/test_hflx_limiter_pd_stencil_01.py new file mode 100644 index 0000000000..a75ffe6bba --- /dev/null +++ b/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -0,0 +1,74 @@ +# 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 + +import numpy as np + +from icon4py.advection.hflx_limiter_pd_stencil_01 import ( + hflx_limiter_pd_stencil_01, +) +from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + +def hflx_limiter_pd_stencil_01_numpy( + c2e: np.array, + geofac_div: np.array, + p_cc: np.array, + p_rhodz_now: np.array, + p_mflx_tracer_h: np.array, + p_dtime, + dbl_eps, +): +# p_m: np.array + p_mflx_tracer_h_c2e = p_mflx_tracer_h[c2e] + #p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) + max(0.0, geofac_div[:, 1]*p_mflx_tracer_h_c2e[:, 1]*p_dtime) + max(0.0, geofac_div[:, 2]*p_mflx_tracer_h_c2e[:, 2]*p_dtime) + geofac_div=np.expand_dims(geofac_div,axis=-1) +# p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) + p_m = np.sum(np.maximum(np.zeros(p_mflx_tracer_h_c2e.shape,dtype=float), p_mflx_tracer_h_c2e*geofac_div*p_dtime)) + r_m = np.minimum(np.ones(p_cc.shape,dtype=float),(p_cc*p_rhodz_now)/(p_m+dbl_eps)) + + return r_m + +def test_hflx_limiter_pd_stencil_01(): + mesh = SimpleMesh() + geofac_div = random_field(mesh, CellDim, C2EDim) + p_cc = random_field(mesh, CellDim, KDim) + p_rhodz_now = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + r_m = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) + + ref = hflx_limiter_pd_stencil_01_numpy( + mesh.c2e, + np.asarray(geofac_div), + np.asarray(p_cc), + np.asarray(p_rhodz_now), + np.asarray(p_mflx_tracer_h), + p_dtime, + dbl_eps, + ) + + hflx_limiter_pd_stencil_01( + geofac_div, + p_cc, + p_rhodz_now, + p_mflx_tracer_h, + r_m, + p_dtime, + dbl_eps, + offset_provider={ + "C2E": mesh.get_c2e_offset_provider(), + }, + ) + assert np.allclose(r_m, ref) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py new file mode 100644 index 0000000000..2079beb600 --- /dev/null +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -0,0 +1,70 @@ +# 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 + +import numpy as np + +from icon4py.advection.hflx_limiter_pd_stencil_02 import ( + hflx_limiter_pd_stencil_02, +) +from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + +def hflx_limiter_pd_stencil_02_numpy( + e2c: np.array, + refin_ctrl: np.array, + r_m: np.array, + p_mflx_tracer_h_in: np.array, + bound, +): + r_m_e2c = r_m[e2c] + refin_ctrl = np.expand_dims(refin_ctrl, axis=-1) + p_mflx_tracer_h_out = np.where ( + refin_ctrl != bound, + np.where ( + p_mflx_tracer_h_in >=0, + p_mflx_tracer_h_in * r_m_e2c[:, 0], + p_mflx_tracer_h_in * r_m_e2c[:, 1], + ), + p_mflx_tracer_h_in, + ) + return p_mflx_tracer_h_out + +def test_hflx_limiter_pd_stencil_02(): + mesh = SimpleMesh() + + refin_ctrl = random_field(mesh, EdgeDim) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + p_mflx_tracer_h_out = random_field(mesh, EdgeDim, KDim) + bound=np.float64(7.0) + + ref = hflx_limiter_pd_stencil_02_numpy( + mesh.e2c, + np.asarray(refin_ctrl), + np.asarray(r_m), + np.asarray(p_mflx_tracer_h_in), + bound, + ) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + p_mflx_tracer_h_out, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_out, ref) diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/advection/tests/test_rbf_intp_edge_stencil_01.py new file mode 100644 index 0000000000..f9a20110e6 --- /dev/null +++ b/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -0,0 +1,57 @@ +# 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 + +import numpy as np + +from icon4py.advection.rbf_intp_edge_stencil_01 import ( + rbf_intp_edge_stencil_01, +) +from icon4py.common.dimension import EdgeDim, E2C2EDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + +def rbf_intp_edge_stencil_01_numpy( + e2c2e: np.array, + p_vn_in: np.array, + ptr_coeff: np.array, +) -> np.array: + + ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) + p_vt_out = np.sum(p_vn_in[e2c2e] * ptr_coeff, axis=1) +# p_vn_ptr = p_vn_in[e2c2e] * ptr_coeff +# p_vt_out = np.sum(p_vn_ptr, axis=1) + + return p_vt_out + +def test_rbf_intp_edge_stencil_01(): + mesh= SimpleMesh() + + p_vn_in = random_field(mesh, EdgeDim, KDim) + ptr_coeff = random_field(mesh, EdgeDim, E2C2EDim) + p_vt_out = zero_field(mesh, EdgeDim, KDim) + + ref = rbf_intp_edge_stencil_01_numpy( + mesh.e2c2e, + np.asarray(p_vn_in), + np.asarray(ptr_coeff), + ) + + rbf_intp_edge_stencil_01( + p_vn_in, + ptr_coeff, + p_vt_out, + offset_provider={ + "E2C2E": mesh.get_e2c2e_offset_provider(), + }, + ) + assert np.allclose(p_vt_out, ref) diff --git a/advection/tests/test_step_advection_stencil_04.py b/advection/tests/test_step_advection_stencil_04.py new file mode 100644 index 0000000000..463293458b --- /dev/null +++ b/advection/tests/test_step_advection_stencil_04.py @@ -0,0 +1,52 @@ +# 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 + +import numpy as np + +from icon4py.advection.step_advection_stencil_04 import ( + step_advection_stencil_04, +) + +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + +def step_advection_stencil_04_numpy( + p_tracer_now: np.array, + p_tracer_new: np.array, + p_dtime, +) -> np.array: + opt_ddt_tracer_adv = (p_tracer_new-p_tracer_now)/p_dtime + return opt_ddt_tracer_adv + +def test_step_advection_stencil_04(): + mesh = SimpleMesh() + + p_tracer_now = random_field(mesh, CellDim, KDim) + p_tracer_new = random_field(mesh, CellDim, KDim) + opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = step_advection_stencil_04_numpy( + np.asarray(p_tracer_now), + np.asarray(p_tracer_new), + p_dtime, + ) + step_advection_stencil_04( + p_tracer_now, + p_tracer_new, + opt_ddt_tracer_adv, + p_dtime, + offset_provider={}, +) + assert np.allclose(opt_ddt_tracer_adv, ref) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py new file mode 100644 index 0000000000..c259bea688 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py @@ -0,0 +1,37 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _compute_airmass( + rho_in: Field[[CellDim, KDim], float], + ddqz_z_full_in: Field[[CellDim, KDim], float], + deepatmo_t1mc_in: Field[[KDim], float], +) -> Field[[CellDim, KDim], float]: + return rho_in * ddqz_z_full_in * deepatmo_t1mc_in + + + +@program +def compute_airmass( + rho_in: Field[[CellDim, KDim], float], + ddqz_z_full_in: Field[[CellDim, KDim], float], + deepatmo_t1mc_in: Field[[KDim], float], + airmass_out: Field[[CellDim, KDim], float], +): + _compute_airmass(rho_in, ddqz_z_full_in, deepatmo_t1mc_in, out=airmass_out) diff --git a/atm_dyn_iconam/tests/test_compute_airmass.py b/atm_dyn_iconam/tests/test_compute_airmass.py new file mode 100644 index 0000000000..0ca950c892 --- /dev/null +++ b/atm_dyn_iconam/tests/test_compute_airmass.py @@ -0,0 +1,44 @@ +# 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 + +import numpy as np + +from icon4py.atm_dyn_iconam.compute_airmass import ( + compute_airmass, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field + + +def compute_airmass_numpy( + rho_in: np.array, ddqz_z_full_in: np.array, deepatmo_t1mc_in: np.array +) -> np.array: + return rho_in * ddqz_z_full_in * deepatmo_t1mc_in + + +def test_compute_airmass(): + mesh = SimpleMesh() + + rho_in = random_field(mesh, CellDim, KDim) + ddqz_z_full_in = random_field(mesh, CellDim, KDim) + deepatmo_t1mc_in = random_field(mesh, KDim) + airmass_out = random_field(mesh, CellDim, KDim) + + ref = compute_airmass_numpy( + np.asarray(rho_in), np.asarray(ddqz_z_full_in), np.asarray(deepatmo_t1mc_in) + ) + compute_airmass( + rho_in, ddqz_z_full_in, deepatmo_t1mc_in, airmass_out, offset_provider={} + ) + assert np.allclose(airmass_out, ref) diff --git a/requirements-dev.txt b/requirements-dev.txt index bacaeba53c..975538dfc1 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,3 +5,4 @@ -e ./pyutils -e ./testutils -e ./atm_dyn_iconam +-e ./advection diff --git a/requirements.txt b/requirements.txt index 4e67c2a537..1c6af97a2e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ ./pyutils ./testutils ./atm_dyn_iconam +./advection From 6fda4466f1979b622013b5bfba77cde2f03b219e Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 20 Sep 2022 16:26:40 +0200 Subject: [PATCH 002/105] Fix hflx limiter pd stencil (#91) * run pre-commit * fix hflx_limiter_pd_stencil_01 --- advection/setup.py | 1 + advection/src/icon4py/advection/__init__.py | 1 - .../advection/hflx_limiter_pd_stencil_01.py | 48 ++++++++++--------- .../advection/hflx_limiter_pd_stencil_02.py | 24 +++++----- .../advection/rbf_intp_edge_stencil_01.py | 8 ++-- .../advection/step_advection_stencil_04.py | 8 +++- .../tests/test_hflx_limiter_pd_stencil_01.py | 26 +++++----- .../tests/test_hflx_limiter_pd_stencil_02.py | 18 +++---- .../tests/test_rbf_intp_edge_stencil_01.py | 19 ++++---- .../tests/test_step_advection_stencil_04.py | 23 +++++---- .../icon4py/atm_dyn_iconam/compute_airmass.py | 1 - atm_dyn_iconam/tests/test_compute_airmass.py | 4 +- 12 files changed, 95 insertions(+), 86 deletions(-) diff --git a/advection/setup.py b/advection/setup.py index ceec936ea5..9c9f7b81c8 100644 --- a/advection/setup.py +++ b/advection/setup.py @@ -1,3 +1,4 @@ +# ICON4Py - ICON inspired code in Python and GT4Py # # Copyright (c) 2022, ETH Zurich and MeteoSwiss # All rights reserved. diff --git a/advection/src/icon4py/advection/__init__.py b/advection/src/icon4py/advection/__init__.py index a6d2d236c9..15dfdb0098 100644 --- a/advection/src/icon4py/advection/__init__.py +++ b/advection/src/icon4py/advection/__init__.py @@ -10,4 +10,3 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py index 94c13720b0..55d84ef3ef 100644 --- a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py @@ -12,45 +12,49 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, maximum, minimum, neighbor_sum, broadcast -from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim, C2E +from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum + +from icon4py.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim + @field_operator def _hflx_limiter_pd_stencil_01( - geofac_div: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CEDim], float], p_cc: Field[[CellDim, KDim], float], p_rhodz_now: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - p_dtime:float, + p_dtime: float, dbl_eps: float, ) -> Field[[CellDim, KDim], float]: -# p_m: Field[[CellDim, KDim], float] -# p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) -# p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[1])*geofac_div(1)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[2])*geofac_div(2)*p_dtime) - #p_m = neighbor_sum(a_a, axis=C2EDim) - #r_m = minimum(1.0, (p_cc*p_rhodz_now)/(p_m+dbl_eps)) -# r_m = minimum(1.0, (p_cc*p_rhodz_now)/(1.0+dbl_eps)) - - p_m = neighbor_sum(maximum(broadcast(0.0, (CellDim, KDim)), p_mflx_tracer_h(C2E)*geofac_div*p_dtime), axis=C2EDim) - r_m = minimum(broadcast(1.0, (CellDim, KDim)), (p_cc*p_rhodz_now)/(p_m+dbl_eps)) + zero = broadcast(0.0, (CellDim, KDim)) + + pm_0 = maximum(zero, p_mflx_tracer_h(C2E[0]) * geofac_div(C2CE[0]) * p_dtime) + pm_1 = maximum(zero, p_mflx_tracer_h(C2E[1]) * geofac_div(C2CE[1]) * p_dtime) + pm_2 = maximum(zero, p_mflx_tracer_h(C2E[2]) * geofac_div(C2CE[2]) * p_dtime) + p_m = pm_0 + pm_1 + pm_2 + r_m = minimum( + broadcast(1.0, (CellDim, KDim)), (p_cc * p_rhodz_now) / (p_m + dbl_eps) + ) + return r_m + @program def hflx_limiter_pd_stencil_01( - geofac_div: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CEDim], float], p_cc: Field[[CellDim, KDim], float], p_rhodz_now: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], r_m: Field[[CellDim, KDim], float], - p_dtime:float, + p_dtime: float, dbl_eps: float, ): _hflx_limiter_pd_stencil_01( - geofac_div, - p_cc, - p_rhodz_now, - p_mflx_tracer_h, - p_dtime, - dbl_eps, - out=r_m, + geofac_div, + p_cc, + p_rhodz_now, + p_mflx_tracer_h, + p_dtime, + dbl_eps, + out=r_m, ) diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py index 2285800f89..ac59f1690f 100644 --- a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py @@ -14,26 +14,28 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, where -from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2C +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim + @field_operator def _hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], float], + refin_ctrl: Field[[EdgeDim], float], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], bound: float, ) -> Field[[EdgeDim, KDim], float]: - p_mflx_tracer_h_out = where ( + p_mflx_tracer_h_out = where( refin_ctrl == bound, p_mflx_tracer_h_in, - where ( - (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in==0.0), + where( + (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in == 0.0), p_mflx_tracer_h_in * r_m(E2C[0]), p_mflx_tracer_h_in * r_m(E2C[1]), ), ) return p_mflx_tracer_h_out + @program def hflx_limiter_pd_stencil_02( refin_ctrl: Field[[EdgeDim], float], @@ -43,9 +45,9 @@ def hflx_limiter_pd_stencil_02( bound: float, ): _hflx_limiter_pd_stencil_02( - refin_ctrl, - r_m, - p_mflx_tracer_h_in, - bound, - out=p_mflx_tracer_h_out, -) + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + out=p_mflx_tracer_h_out, + ) diff --git a/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py b/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py index 76a97b58c9..46e1069428 100644 --- a/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py +++ b/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py @@ -19,17 +19,17 @@ @field_operator def _rbf_intp_edge_stencil_01( - p_vn_in: Field[[EdgeDim, KDim], float], - # ptr_coeff: Field[[EdgeDim, E2CDim, E2C2EDim], float] + p_vn_in: Field[[EdgeDim, KDim], float], ptr_coeff: Field[[EdgeDim, E2C2EDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim) + p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim) return p_vt_out + @program def rbf_intp_edge_stencil_01( p_vn_in: Field[[EdgeDim, KDim], float], ptr_coeff: Field[[EdgeDim, E2C2EDim], float], p_vt_out: Field[[EdgeDim, KDim], float], ): - _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out) + _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out) diff --git a/advection/src/icon4py/advection/step_advection_stencil_04.py b/advection/src/icon4py/advection/step_advection_stencil_04.py index f299d11d98..f0177fabef 100644 --- a/advection/src/icon4py/advection/step_advection_stencil_04.py +++ b/advection/src/icon4py/advection/step_advection_stencil_04.py @@ -16,15 +16,17 @@ from icon4py.common.dimension import CellDim, KDim + @field_operator def _step_advection_stencil_04( p_tracer_now: Field[[CellDim, KDim], float], p_tracer_new: Field[[CellDim, KDim], float], p_dtime: float, ) -> Field[[CellDim, KDim], float]: - opt_ddt_tracer_adv=(p_tracer_new-p_tracer_now)/p_dtime + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime return opt_ddt_tracer_adv + @program def step_advection_stencil_04( p_tracer_now: Field[[CellDim, KDim], float], @@ -32,4 +34,6 @@ def step_advection_stencil_04( opt_ddt_tracer_adv: Field[[CellDim, KDim], float], p_dtime: float, ): - _step_advection_stencil_04(p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv) + _step_advection_stencil_04( + p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv + ) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/advection/tests/test_hflx_limiter_pd_stencil_01.py index a75ffe6bba..6e7d85da9e 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -11,14 +11,16 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np +import numpy as np +from functional.iterator.embedded import StridedNeighborOffsetProvider from icon4py.advection.hflx_limiter_pd_stencil_01 import ( hflx_limiter_pd_stencil_01, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim +from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field + def hflx_limiter_pd_stencil_01_numpy( c2e: np.array, @@ -29,16 +31,17 @@ def hflx_limiter_pd_stencil_01_numpy( p_dtime, dbl_eps, ): -# p_m: np.array - p_mflx_tracer_h_c2e = p_mflx_tracer_h[c2e] - #p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) + max(0.0, geofac_div[:, 1]*p_mflx_tracer_h_c2e[:, 1]*p_dtime) + max(0.0, geofac_div[:, 2]*p_mflx_tracer_h_c2e[:, 2]*p_dtime) - geofac_div=np.expand_dims(geofac_div,axis=-1) -# p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) - p_m = np.sum(np.maximum(np.zeros(p_mflx_tracer_h_c2e.shape,dtype=float), p_mflx_tracer_h_c2e*geofac_div*p_dtime)) - r_m = np.minimum(np.ones(p_cc.shape,dtype=float),(p_cc*p_rhodz_now)/(p_m+dbl_eps)) + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 0]] * geofac_div[:, 0] * p_dtime) + p_m_1 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 1]] * geofac_div[:, 1] * p_dtime) + p_m_2 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 2]] * geofac_div[:, 2] * p_dtime) + + p_m = p_m_0 + p_m_1 + p_m_2 + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) return r_m + def test_hflx_limiter_pd_stencil_01(): mesh = SimpleMesh() geofac_div = random_field(mesh, CellDim, C2EDim) @@ -60,7 +63,7 @@ def test_hflx_limiter_pd_stencil_01(): ) hflx_limiter_pd_stencil_01( - geofac_div, + as_1D_sparse_field(geofac_div, CEDim), p_cc, p_rhodz_now, p_mflx_tracer_h, @@ -68,6 +71,7 @@ def test_hflx_limiter_pd_stencil_01(): p_dtime, dbl_eps, offset_provider={ + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), "C2E": mesh.get_c2e_offset_provider(), }, ) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index 2079beb600..fdb93587cb 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -11,28 +11,29 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np +import numpy as np from icon4py.advection.hflx_limiter_pd_stencil_02 import ( hflx_limiter_pd_stencil_02, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import random_field + def hflx_limiter_pd_stencil_02_numpy( e2c: np.array, - refin_ctrl: np.array, + refin_ctrl: np.array, r_m: np.array, p_mflx_tracer_h_in: np.array, bound, ): r_m_e2c = r_m[e2c] refin_ctrl = np.expand_dims(refin_ctrl, axis=-1) - p_mflx_tracer_h_out = np.where ( + p_mflx_tracer_h_out = np.where( refin_ctrl != bound, - np.where ( - p_mflx_tracer_h_in >=0, + np.where( + p_mflx_tracer_h_in >= 0, p_mflx_tracer_h_in * r_m_e2c[:, 0], p_mflx_tracer_h_in * r_m_e2c[:, 1], ), @@ -40,6 +41,7 @@ def hflx_limiter_pd_stencil_02_numpy( ) return p_mflx_tracer_h_out + def test_hflx_limiter_pd_stencil_02(): mesh = SimpleMesh() @@ -47,7 +49,7 @@ def test_hflx_limiter_pd_stencil_02(): r_m = random_field(mesh, CellDim, KDim) p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) p_mflx_tracer_h_out = random_field(mesh, EdgeDim, KDim) - bound=np.float64(7.0) + bound = np.float64(7.0) ref = hflx_limiter_pd_stencil_02_numpy( mesh.e2c, @@ -58,7 +60,7 @@ def test_hflx_limiter_pd_stencil_02(): ) hflx_limiter_pd_stencil_02( - refin_ctrl, + refin_ctrl, r_m, p_mflx_tracer_h_in, p_mflx_tracer_h_out, diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/advection/tests/test_rbf_intp_edge_stencil_01.py index f9a20110e6..84bd29aa43 100644 --- a/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -11,34 +11,31 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np +import numpy as np -from icon4py.advection.rbf_intp_edge_stencil_01 import ( - rbf_intp_edge_stencil_01, -) -from icon4py.common.dimension import EdgeDim, E2C2EDim, KDim +from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 +from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field + def rbf_intp_edge_stencil_01_numpy( e2c2e: np.array, - p_vn_in: np.array, + p_vn_in: np.array, ptr_coeff: np.array, ) -> np.array: ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) p_vt_out = np.sum(p_vn_in[e2c2e] * ptr_coeff, axis=1) -# p_vn_ptr = p_vn_in[e2c2e] * ptr_coeff -# p_vt_out = np.sum(p_vn_ptr, axis=1) - return p_vt_out + def test_rbf_intp_edge_stencil_01(): - mesh= SimpleMesh() + mesh = SimpleMesh() p_vn_in = random_field(mesh, EdgeDim, KDim) ptr_coeff = random_field(mesh, EdgeDim, E2C2EDim) - p_vt_out = zero_field(mesh, EdgeDim, KDim) + p_vt_out = zero_field(mesh, EdgeDim, KDim) ref = rbf_intp_edge_stencil_01_numpy( mesh.e2c2e, diff --git a/advection/tests/test_step_advection_stencil_04.py b/advection/tests/test_step_advection_stencil_04.py index 463293458b..6da3ec9b54 100644 --- a/advection/tests/test_step_advection_stencil_04.py +++ b/advection/tests/test_step_advection_stencil_04.py @@ -13,34 +13,33 @@ import numpy as np -from icon4py.advection.step_advection_stencil_04 import ( - step_advection_stencil_04, -) - +from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.common.dimension import CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field + def step_advection_stencil_04_numpy( - p_tracer_now: np.array, - p_tracer_new: np.array, + p_tracer_now: np.array, + p_tracer_new: np.array, p_dtime, ) -> np.array: - opt_ddt_tracer_adv = (p_tracer_new-p_tracer_now)/p_dtime + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime return opt_ddt_tracer_adv + def test_step_advection_stencil_04(): mesh = SimpleMesh() p_tracer_now = random_field(mesh, CellDim, KDim) p_tracer_new = random_field(mesh, CellDim, KDim) - opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim) - p_dtime = np.float64(5.0) + opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) ref = step_advection_stencil_04_numpy( np.asarray(p_tracer_now), - np.asarray(p_tracer_new), - p_dtime, + np.asarray(p_tracer_new), + p_dtime, ) step_advection_stencil_04( p_tracer_now, @@ -48,5 +47,5 @@ def test_step_advection_stencil_04(): opt_ddt_tracer_adv, p_dtime, offset_provider={}, -) + ) assert np.allclose(opt_ddt_tracer_adv, ref) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py index c259bea688..bac9061ce2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py @@ -26,7 +26,6 @@ def _compute_airmass( return rho_in * ddqz_z_full_in * deepatmo_t1mc_in - @program def compute_airmass( rho_in: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/tests/test_compute_airmass.py b/atm_dyn_iconam/tests/test_compute_airmass.py index 0ca950c892..881838f8b5 100644 --- a/atm_dyn_iconam/tests/test_compute_airmass.py +++ b/atm_dyn_iconam/tests/test_compute_airmass.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.atm_dyn_iconam.compute_airmass import ( - compute_airmass, -) +from icon4py.atm_dyn_iconam.compute_airmass import compute_airmass from icon4py.common.dimension import CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field From 7cd58eddfa5e0650c061584762746a121e234f36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 27 Sep 2022 15:37:25 +0200 Subject: [PATCH 003/105] Fixes to hflx_limiter_pd_stencil_02 to make it pass probtest --- .../src/icon4py/advection/hflx_limiter_pd_stencil_02.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py index ac59f1690f..857ee7584d 100644 --- a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py @@ -40,14 +40,13 @@ def _hflx_limiter_pd_stencil_02( def hflx_limiter_pd_stencil_02( refin_ctrl: Field[[EdgeDim], float], r_m: Field[[CellDim, KDim], float], - p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], - p_mflx_tracer_h_out: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], bound: float, ): _hflx_limiter_pd_stencil_02( refin_ctrl, r_m, - p_mflx_tracer_h_in, + p_mflx_tracer_h, bound, - out=p_mflx_tracer_h_out, + out=p_mflx_tracer_h, ) From 8b324a3bec04f3154d183e81d6f6e92133380133 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 12 Oct 2022 14:00:12 +0200 Subject: [PATCH 004/105] Tracer advection stencil porting ml -> tracer advection stencils (#98) * adding hflx_limiter_mo_stencil_02 * changes order of args and return a tuple * Revert "changes order of args and return a tuple" This reverts commit 99638df7024fcbcccf1ad37cda2b5e7651efb918. * porting step_advection_stencil_01 * porting step_advection_stencil_02 * porting upwind_hflux_miura_stencil_02 * hflx_limiter_mo_stencil_04 * fix step_advection_stencil_01 * hflx_limiter_mo_stencil_03 * clean up: move upwind_hflux_miura_stencil_02.py * add test for conditional branch in hflx_limiter_mo_stencil_04 * switch to using min_over in hflx_limiter_mo_stencil_03 * refactor hflx_limiter_mo_stencil_04, fix tests for Koffset * add upwind_vflux_ppm_stencil_01 * upwind_hflux_miura_stencil_02 - fix sparse field issue, split into separate sums instead of using neighbor_sum * remove commented out code * fix: upwind_vflux_ppm_stencil_01 * fix: hflx_limiter_mo_stencil_01: return tuple, use int32 * hflx_limiter_pd_stencil_02.py: use int fields --- .../advection/hflx_limiter_mo_stencil_02.py | 59 ++++++++ .../advection/hflx_limiter_mo_stencil_03.py | 92 +++++++++++++ .../advection/hflx_limiter_mo_stencil_04.py | 44 ++++++ .../advection/hflx_limiter_pd_stencil_02.py | 12 +- .../advection/step_advection_stencil_01.py | 50 +++++++ .../advection/step_advection_stencil_02.py | 50 +++++++ .../upwind_hflux_miura_stencil_02.py | 49 +++++++ .../advection/upwind_vflux_ppm_stencil_01.py | 40 ++++++ .../tests/test_hflx_limiter_mo_stencil_02.py | 103 ++++++++++++++ .../tests/test_hflx_limiter_mo_stencil_03.py | 126 ++++++++++++++++++ .../tests/test_hflx_limiter_mo_stencil_04.py | 61 +++++++++ .../tests/test_hflx_limiter_pd_stencil_02.py | 71 ++++++++-- .../tests/test_step_advection_stencil_01.py | 63 +++++++++ .../tests/test_step_advection_stencil_02.py | 64 +++++++++ .../test_upwind_hflux_miura_stencil_02.py | 72 ++++++++++ .../tests/test_upwind_vflux_ppm_stencil_01.py | 49 +++++++ common/src/icon4py/common/dimension.py | 3 + 17 files changed, 994 insertions(+), 14 deletions(-) create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py create mode 100644 advection/src/icon4py/advection/step_advection_stencil_01.py create mode 100644 advection/src/icon4py/advection/step_advection_stencil_02.py create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py create mode 100644 advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py create mode 100644 advection/tests/test_hflx_limiter_mo_stencil_02.py create mode 100644 advection/tests/test_hflx_limiter_mo_stencil_03.py create mode 100644 advection/tests/test_hflx_limiter_mo_stencil_04.py create mode 100644 advection/tests/test_step_advection_stencil_01.py create mode 100644 advection/tests/test_step_advection_stencil_02.py create mode 100644 advection/tests/test_upwind_hflux_miura_stencil_02.py create mode 100644 advection/tests/test_upwind_vflux_ppm_stencil_01.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py new file mode 100644 index 0000000000..c9c359621f --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py @@ -0,0 +1,59 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import int32, maximum, minimum, where + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound) + z_tracer_new_tmp = where( + condition, + minimum(1.1 * p_cc, maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + return where( + condition, + (maximum(p_cc, z_tracer_new_tmp), minimum(p_cc, z_tracer_new_tmp)), + (z_tracer_new_low, z_tracer_new_low), + ) + + +@program +def hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + out=(z_tracer_max, z_tracer_min), + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py new file mode 100644 index 0000000000..5969a171fd --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py @@ -0,0 +1,92 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import max_over, maximum, min_over, minimum + +from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_max = beta_fct * maximum( + max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max + ) + z_min = r_beta_fct * minimum( + min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min + ) + return z_max, z_min + + +@program +def hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_max, z_min) + ) + + +@field_operator +def _hflx_limiter_mo_stencil_03( + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], + dbl_eps: float, +) -> tuple[[Field[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + return r_p, r_m + + +@program +def hflx_limiter_mo_stencil_03( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + dbl_eps: float, + r_p: Field[[CellDim, KDim], float], + r_m: Field[[CellDim, KDim], float], +): + + _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_min, z_max) + ) + _hflx_limiter_mo_stencil_03( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_max, + z_min, + dbl_eps, + out=(r_p, r_m), + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py new file mode 100644 index 0000000000..8a97740918 --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py @@ -0,0 +1,44 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import minimum, where + +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + r_frac = where( + z_anti >= 0.0, + minimum(r_m(E2C[0]), r_p(E2C[1])), + minimum(r_m(E2C[1]), r_p(E2C[0])), + ) + return z_mflx_low + minimum(1.0, r_frac) * z_anti + + +@program +def hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_04(z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h) diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py index 857ee7584d..5d4d239f08 100644 --- a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py @@ -12,23 +12,23 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where +from functional.ffront.fbuiltins import Field, int32, where from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator def _hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], float], + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], - bound: float, + bound: int32, ) -> Field[[EdgeDim, KDim], float]: p_mflx_tracer_h_out = where( refin_ctrl == bound, p_mflx_tracer_h_in, where( - (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in == 0.0), + p_mflx_tracer_h_in >= 0.0, p_mflx_tracer_h_in * r_m(E2C[0]), p_mflx_tracer_h_in * r_m(E2C[1]), ), @@ -38,10 +38,10 @@ def _hflx_limiter_pd_stencil_02( @program def hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], float], + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: float, + bound: int32, ): _hflx_limiter_pd_stencil_02( refin_ctrl, diff --git a/advection/src/icon4py/advection/step_advection_stencil_01.py b/advection/src/icon4py/advection/step_advection_stencil_01.py new file mode 100644 index 0000000000..0a2ac7e3c8 --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_01.py @@ -0,0 +1,50 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + k_offset_up_low = p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + return rhodz_ast + k_offset_up_low + + +@program +def step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_01( + rhodz_ast, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/advection/src/icon4py/advection/step_advection_stencil_02.py b/advection/src/icon4py/advection/step_advection_stencil_02.py new file mode 100644 index 0000000000..e004903f5a --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_02.py @@ -0,0 +1,50 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import maximum + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + return maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + + +@program +def step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_02( + p_rhodz_new, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py new file mode 100644 index 0000000000..f1c8abca13 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py @@ -0,0 +1,49 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim + + +@field_operator +def _upwind_hflux_miura_stencil_02( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + p_coeff_1 = ( + lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + p_coeff_2 = ( + lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + return p_coeff_1, p_coeff_2 + + +@program +def upwind_hflux_miura_stencil_02( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], + p_coeff_1: Field[[CellDim, KDim], float], + p_coeff_2: Field[[CellDim, KDim], float], +): + _upwind_hflux_miura_stencil_02( + p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2) + ) diff --git a/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py b/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py new file mode 100644 index 0000000000..6a027a7482 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,40 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + + return z_delta_q, z_a1 + + +@program +def upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_delta_q: Field[[CellDim, KDim], float], + z_a1: Field[[CellDim, KDim], float], +): + _upwind_vflux_ppm_stencil_01(z_face_up, z_face_low, p_cc, out=(z_delta_q, z_a1)) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py new file mode 100644 index 0000000000..13ed590a59 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -0,0 +1,103 @@ +# 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 + +import numpy as np +from numpy import int32 + +from icon4py.advection.hflx_limiter_mo_stencil_02 import ( + hflx_limiter_mo_stencil_02, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import constant_field, random_field, zero_field + + +def hflx_limiter_mo_stencil_02_numpy( + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + lo_bound: float, + hi_bound: float, +): + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), + ) + z_tracer_new_tmp = np.where( + condition, + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + z_tracer_max = np.where( + condition, np.maximum(p_cc, z_tracer_new_tmp), z_tracer_new_low + ) + z_tracer_min = np.where( + condition, np.minimum(p_cc, z_tracer_new_tmp), z_tracer_new_low + ) + return z_tracer_max, z_tracer_min + + +def test_hflx_limiter_mo_stencil_02_some_matching_condition(): + mesh = SimpleMesh() + hi_bound = np.int32(1) + lo_bound = np.int32(5) + refin_ctrl = constant_field(mesh, hi_bound, CellDim, dtype=int32) + refin_ctrl[0:2] = np.int32(3) + p_cc = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( + np.asarray(refin_ctrl), + np.asarray(p_cc), + np.asarray(z_tracer_new_low), + lo_bound, + hi_bound, + ) + + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + assert np.allclose(z_tracer_max, ref_max) + assert np.allclose(z_tracer_min, ref_min) + + +def test_hflx_limiter_mo_stencil_02_none_matching_condition(): + mesh = SimpleMesh() + hi_bound = np.int32(3) + lo_bound = np.int32(1) + refin_ctrl = constant_field(mesh, 2, CellDim, dtype=int32) + p_cc = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + assert np.allclose(z_tracer_min, z_tracer_new_low) + assert np.allclose(z_tracer_max, z_tracer_new_low) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py new file mode 100644 index 0000000000..f696b358b3 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -0,0 +1,126 @@ +# 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 +import numpy as np + +from icon4py.advection.hflx_limiter_mo_stencil_03 import ( + hflx_limiter_mo_stencil_03, + hflx_limiter_mo_stencil_03_min_max, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def hflx_limiter_mo_stencil_03_numpy( + c2e2c: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: np.ndarray, + z_mflx_anti_out: np.ndarray, + z_tracer_new_low: np.ndarray, + dbl_eps: float, +): + z_max, z_min = hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out * dbl_eps) + return r_p, r_m + + +def hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c: np.array, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, +) -> tuple[np.ndarray]: + z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) + z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) + return z_max, z_min + + +def test_hflx_diffusion_mo_stencil_03_min_max(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + z_max = zero_field(mesh, CellDim, KDim) + z_min = zero_field(mesh, CellDim, KDim) + beta_fct = 0.9 + r_beta_fct = 0.3 + z_max_ref, z_min_ref = hflx_limiter_mo_stencil_03_min_max_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + ) + hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_max, + z_min, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + assert np.allclose(z_max, z_max_ref) + assert np.allclose(z_min, z_min_ref) + + +def test_hflx_diffusion_mo_stencil_03(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + z_max = zero_field(mesh, CellDim, KDim) + z_min = zero_field(mesh, CellDim, KDim) + beta_fct = 0.4 + r_beta_fct = 0.6 + z_mflx_anti_in = random_field(mesh, CellDim, KDim) + z_mflx_anti_out = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + dbl_eps = 1e-5 + r_p = zero_field(mesh, CellDim, KDim) + r_m = zero_field(mesh, CellDim, KDim) + + r_p_ref, r_m_ref = hflx_limiter_mo_stencil_03_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + np.asarray(z_mflx_anti_in), + np.asarray(z_mflx_anti_out), + np.asarray(z_tracer_new_low), + dbl_eps, + ) + + hflx_limiter_mo_stencil_03( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_max, + z_min, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + dbl_eps, + r_p, + r_m, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + np.allclose(r_p_ref, r_p) + np.allclose(r_m_ref, r_m) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/advection/tests/test_hflx_limiter_mo_stencil_04.py new file mode 100644 index 0000000000..bdb7211b88 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -0,0 +1,61 @@ +# 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 + +import numpy as np + +from icon4py.advection.hflx_limiter_mo_stencil_04 import ( + hflx_limiter_mo_stencil_04, +) +from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def hflx_limiter_mo_stencil_04_numpy( + e2c: np.ndarray, + z_anti: np.ndarray, + r_m: np.ndarray, + r_p: np.ndarray, + z_mflx_low: np.ndarray, +): + r_frac = np.where( + z_anti >= 0, + np.minimum(r_m[e2c[:, 0]], r_p[e2c[:, 1]]), + np.minimum(r_m[e2c[:, 1]], r_p[e2c[:, 0]]), + ) + return z_mflx_low + np.minimum(1.0, r_frac) * z_anti + + +def test_hflx_limiter_mo_stencil_04(): + mesh = SimpleMesh() + z_anti = random_field(mesh, EdgeDim, KDim, low=-2.0, high=2.0) + r_m = random_field(mesh, CellDim, KDim) + r_p = random_field(mesh, CellDim, KDim) + z_mflx_low = random_field(mesh, EdgeDim, KDim) + p_mflx_tracer_h = zero_field(mesh, EdgeDim, KDim) + ref = hflx_limiter_mo_stencil_04_numpy( + mesh.e2c, + np.asarray(z_anti), + np.asarray(r_m), + np.asarray(r_p), + np.asarray(z_mflx_low), + ) + hflx_limiter_mo_stencil_04( + z_anti, + r_m, + r_p, + z_mflx_low, + p_mflx_tracer_h, + offset_provider={"E2C": mesh.get_e2c_offset_provider()}, + ) + assert np.allclose(p_mflx_tracer_h, ref) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index fdb93587cb..57eb18aa59 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field +from icon4py.testutils.utils import constant_field, random_field def hflx_limiter_pd_stencil_02_numpy( @@ -42,14 +42,12 @@ def hflx_limiter_pd_stencil_02_numpy( return p_mflx_tracer_h_out -def test_hflx_limiter_pd_stencil_02(): +def test_hflx_limiter_pd_stencil_02_nowhere_matching_refin_ctl(): mesh = SimpleMesh() - - refin_ctrl = random_field(mesh, EdgeDim) + bound = np.int32(7) + refin_ctrl = constant_field(mesh, 4, EdgeDim, dtype=np.int32) r_m = random_field(mesh, CellDim, KDim) p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) - p_mflx_tracer_h_out = random_field(mesh, EdgeDim, KDim) - bound = np.float64(7.0) ref = hflx_limiter_pd_stencil_02_numpy( mesh.e2c, @@ -63,10 +61,67 @@ def test_hflx_limiter_pd_stencil_02(): refin_ctrl, r_m, p_mflx_tracer_h_in, - p_mflx_tracer_h_out, bound, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), }, ) - assert np.allclose(p_mflx_tracer_h_out, ref) + assert np.allclose(p_mflx_tracer_h_in, ref) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(4) + refin_ctrl = constant_field(mesh, 5, EdgeDim, dtype=np.int32) + refin_ctrl[2:6] = bound + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl_does_not_change_inout_arg(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py new file mode 100644 index 0000000000..fc2290de19 --- /dev/null +++ b/advection/tests/test_step_advection_stencil_01.py @@ -0,0 +1,63 @@ +# 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 + +import numpy as np + +from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def step_advection_stencil_01_numpy( + rhodz_ast: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + p_mlfx_contra_v1 = np.roll(p_mflx_contra_v, axis=1, shift=-1) + tmp = pd_time * ( + p_mlfx_contra_v1 * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + return rhodz_ast + tmp + + +def test_step_advection_stencil_01(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_01_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_01( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py new file mode 100644 index 0000000000..a9a5e99ca9 --- /dev/null +++ b/advection/tests/test_step_advection_stencil_02.py @@ -0,0 +1,64 @@ +# 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 + +import numpy as np + +from icon4py.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def step_advection_stencil_02_numpy( + rhodz_new: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + + tmp = ( + np.roll(p_mflx_contra_v, axis=1, shift=-1) * deepatmo_divzl + - p_mflx_contra_v * deepatmo_divzu + ) + return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp + + +def test_step_advection_stencil_02(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_02_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_02( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py new file mode 100644 index 0000000000..3dae540683 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -0,0 +1,72 @@ +# 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 + +import numpy as np +from functional.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.upwind_hflux_miura_stencil_02 import ( + upwind_hflux_miura_stencil_02, +) +from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field + + +def upwind_hflux_miura_stencil_02_numpy( + c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_pseudoinv_1: np.ndarray, + lsq_pseudoinv_2: np.ndarray, +) -> tuple[np.ndarray]: + p_cc_e = np.expand_dims(p_cc, axis=1) + n_diff = p_cc[c2e2c] - p_cc_e + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + p_coeff_1 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_2 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return p_coeff_1, p_coeff_2 + + +def test_upwind_hflux_miura_stencil_02(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECDim) + + lsq_pseudoinv_2 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECDim) + p_coeff_1 = zero_field(mesh, CellDim, KDim) + p_coeff_2 = zero_field(mesh, CellDim, KDim) + + ref_1, ref_2 = upwind_hflux_miura_stencil_02_numpy( + mesh.c2e2c, + np.asarray(p_cc), + np.asarray(lsq_pseudoinv_1), + np.asarray(lsq_pseudoinv_2), + ) + + upwind_hflux_miura_stencil_02( + p_cc, + lsq_pseudoinv_1_field, + lsq_pseudoinv_2_field, + p_coeff_1, + p_coeff_2, + offset_provider={ + "C2E2C": mesh.get_c2e2c_offset_provider(), + "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), + }, + ) + co1 = np.asarray(p_coeff_1) + co2 = np.asarray(p_coeff_2) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/advection/tests/test_upwind_vflux_ppm_stencil_01.py new file mode 100644 index 0000000000..36f0d8246e --- /dev/null +++ b/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np + +from icon4py.advection.upwind_vflux_ppm_stencil_01 import ( + upwind_vflux_ppm_stencil_01, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def upwind_vflux_ppm_stencil_01_numpy( + z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray +) -> tuple[np.ndarray]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + return z_delta_q, z_a1 + + +def test_upwind_vflux_ppm_stencil_01(): + mesh = SimpleMesh() + z_face_up = random_field(mesh, CellDim, KDim) + z_face_down = random_field(mesh, CellDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + z_delta_q = zero_field(mesh, CellDim, KDim) + z_a1 = zero_field(mesh, CellDim, KDim) + + ref_z_delta_q, ref_z_a1 = upwind_vflux_ppm_stencil_01_numpy( + np.asarray(z_face_up), np.asarray(z_face_down), np.asarray(p_cc) + ) + + upwind_vflux_ppm_stencil_01( + z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={} + ) + + assert np.allclose(ref_z_delta_q, z_delta_q) + assert np.allclose(ref_z_a1, z_a1) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 499bd5ac0d..53f7904193 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -22,6 +22,7 @@ CEDim = Dimension("CE") ECDim = Dimension("EC") ECVDim = Dimension("ECV") +CECDim = Dimension("CEC") E2CDim = Dimension("E2C", DimensionKind.LOCAL) E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) @@ -45,4 +46,6 @@ E2C2EO = FieldOffset("E2C2EO", source=EdgeDim, target=(EdgeDim, E2C2EODim)) E2C2E = FieldOffset("E2C2E", source=EdgeDim, target=(EdgeDim, E2C2EDim)) C2E2C = FieldOffset("C2E2C", source=CellDim, target=(CellDim, C2E2CDim)) +C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim)) + Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) From 24462058334fe7e41360f8de7d88dc4a60c68c94 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 14 Oct 2022 10:14:10 +0200 Subject: [PATCH 005/105] fix return type of hflx_limiter_mo_stencil_03.py (#104) --- advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py index 5969a171fd..19a4ab7e71 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py @@ -56,7 +56,7 @@ def _hflx_limiter_mo_stencil_03( z_max: Field[[CellDim, KDim], float], z_min: Field[[CellDim, KDim], float], dbl_eps: float, -) -> tuple[[Field[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) return r_p, r_m From e7de05ae8684b6d530c1e0fc9105a2db27fbc22c Mon Sep 17 00:00:00 2001 From: abishekg7 <56273301+abishekg7@users.noreply.github.com> Date: Fri, 14 Oct 2022 11:05:23 +0200 Subject: [PATCH 006/105] Tracer advection stencils - agopal (#93) Co-authored-by: Magdalena Luz Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> --- .../advection/face_val_ppm_stencil_05.py | 67 +++++++++++++ .../icon4py/advection/hor_adv_stencil_01.py | 68 ++++++++++++++ .../advection/step_advection_stencil_03.py | 39 ++++++++ .../upwind_hflux_miura_stencil_01.py | 70 ++++++++++++++ .../advection/v_limit_prbl_sm_stencil_01.py | 70 ++++++++++++++ .../icon4py/advection/vert_adv_stencil_01.py | 59 ++++++++++++ .../tests/test_face_val_ppm_stencil_05.py | 83 ++++++++++++++++ advection/tests/test_hor_adv_stencil_01.py | 80 ++++++++++++++++ .../tests/test_step_advection_stencil_03.py | 55 +++++++++++ .../test_upwind_hflux_miura_stencil_01.py | 94 +++++++++++++++++++ advection/tests/test_vert_adv_stencil_01.py | 74 +++++++++++++++ .../tests/test_vlimit_prbl_sm_stencil_01.py | 71 ++++++++++++++ 12 files changed, 830 insertions(+) create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_05.py create mode 100644 advection/src/icon4py/advection/hor_adv_stencil_01.py create mode 100644 advection/src/icon4py/advection/step_advection_stencil_03.py create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py create mode 100644 advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py create mode 100644 advection/src/icon4py/advection/vert_adv_stencil_01.py create mode 100644 advection/tests/test_face_val_ppm_stencil_05.py create mode 100644 advection/tests/test_hor_adv_stencil_01.py create mode 100644 advection/tests/test_step_advection_stencil_03.py create mode 100644 advection/tests/test_upwind_hflux_miura_stencil_01.py create mode 100644 advection/tests/test_vert_adv_stencil_01.py create mode 100644 advection/tests/test_vlimit_prbl_sm_stencil_01.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_05.py b/advection/src/icon4py/advection/face_val_ppm_stencil_05.py new file mode 100644 index 0000000000..7f2f91cd79 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_05.py @@ -0,0 +1,67 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_05( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zgeo1 = p_cellhgt_mc_now(Koff[-1]) / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now(Koff[-2]) + + p_cellhgt_mc_now(Koff[-1]) + + p_cellhgt_mc_now + + p_cellhgt_mc_now(Koff[1]) + ) + zgeo3 = (p_cellhgt_mc_now(Koff[-2]) + p_cellhgt_mc_now(Koff[-1])) / ( + 2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now(Koff[1]) + p_cellhgt_mc_now) / ( + 2.0 * p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1]) + ) + + p_face = ( + p_cc(Koff[-1]) + + zgeo1 * (p_cc - p_cc(Koff[-1])) + + zgeo2 + * ( + (2.0 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc(Koff[-1])) + - zgeo3 * p_cellhgt_mc_now(Koff[-1]) * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope(Koff[-1]) + ) + ) + + return p_face + + +@program +def face_val_ppm_stencil_05( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_05( + p_cc, + p_cellhgt_mc_now, + z_slope, + out=p_face, + ) diff --git a/advection/src/icon4py/advection/hor_adv_stencil_01.py b/advection/src/icon4py/advection/hor_adv_stencil_01.py new file mode 100644 index 0000000000..2cb692a100 --- /dev/null +++ b/advection/src/icon4py/advection/hor_adv_stencil_01.py @@ -0,0 +1,68 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import ( + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) + + +@field_operator +def _hor_adv_stencil_01( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + deepatmo_divh: Field[[KDim], float], + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + tracer_new = ( + tracer_now * rhodz_now + - p_dtime + * deepatmo_divh + * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim) + ) / rhodz_new + + return tracer_new + + +@program +def hor_adv_stencil_01( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + deepatmo_divh: Field[[KDim], float], + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _hor_adv_stencil_01( + p_mflx_tracer_h, + deepatmo_divh, + tracer_now, + rhodz_now, + rhodz_new, + geofac_div, + p_dtime, + out=tracer_new, + ) diff --git a/advection/src/icon4py/advection/step_advection_stencil_03.py b/advection/src/icon4py/advection/step_advection_stencil_03.py new file mode 100644 index 0000000000..2c3728ca0e --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_03.py @@ -0,0 +1,39 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, maximum + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _step_advection_stencil_03( + p_tracer_now: Field[[CellDim, KDim], float], + p_grf_tend_tracer: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + p_tracer_new = maximum(0.0, p_tracer_now + p_dtime * p_grf_tend_tracer) + return p_tracer_new + + +@program +def step_advection_stencil_03( + p_tracer_now: Field[[CellDim, KDim], float], + p_grf_tend_tracer: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _step_advection_stencil_03( + p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py new file mode 100644 index 0000000000..31c44a66c8 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py @@ -0,0 +1,70 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where + +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_stencil_01( + p_vn: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + z_grad_1: Field[[CellDim, KDim], float], + z_grad_2: Field[[CellDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e = where( + p_vn > 0.0, + ( + p_cc(E2C[0]) + + distv_bary_1 * z_grad_1(E2C[0]) + + distv_bary_2 * z_grad_2(E2C[0]) + ) + * p_mass_flx_e, + ( + p_cc(E2C[1]) + + distv_bary_1 * z_grad_1(E2C[1]) + + distv_bary_2 * z_grad_2(E2C[1]) + ) + * p_mass_flx_e, + ) + + return p_out_e + + +@program +def upwind_hflux_miura_stencil_01( + p_vn: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + z_grad_1: Field[[CellDim, KDim], float], + z_grad_2: Field[[CellDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_stencil_01( + p_vn, + p_cc, + distv_bary_1, + distv_bary_2, + z_grad_1, + z_grad_2, + p_mass_flx_e, + out=p_out_e, + ) diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py new file mode 100644 index 0000000000..5d06900ec9 --- /dev/null +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py @@ -0,0 +1,70 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import ( + Field, + FieldOffset, + abs, + broadcast, + minimum, + where, +) + +from icon4py.common.dimension import CellDim, KDim + + +Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) + + +@field_operator +def _v_limit_prbl_sm_stencil_01( + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + z_delta = p_face - p_face(Koff[1]) + z_a6i = -6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) + + q_face_up = broadcast(1.0, (CellDim, KDim)) + q_face_low = broadcast(1.0, (CellDim, KDim)) + + q_face_up, q_face_low = where( + abs(z_delta) < z_a6i, + where( + (p_cc < minimum(p_face, p_face(Koff[1]))), + (p_cc, p_cc), + where( + p_face > p_face(Koff[1]), + (3.0 * p_cc - 2.0 * p_face(Koff[1]), p_face(Koff[1])), + (p_face, 3.0 * p_cc - 2.0 * p_face), + ), + ), + (p_face, p_face(Koff[1])), + ) + + return q_face_up, q_face_low + + +@program +def v_limit_prbl_sm_stencil_01( + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + p_face_up: Field[[CellDim, KDim], float], + p_face_low: Field[[CellDim, KDim], float], +): + _v_limit_prbl_sm_stencil_01( + p_face, + p_cc, + out=(p_face_up, p_face_low), + ) diff --git a/advection/src/icon4py/advection/vert_adv_stencil_01.py b/advection/src/icon4py/advection/vert_adv_stencil_01.py new file mode 100644 index 0000000000..b4a87dd67e --- /dev/null +++ b/advection/src/icon4py/advection/vert_adv_stencil_01.py @@ -0,0 +1,59 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _vert_adv_stencil_01( + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + tracer_new = ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + ) / rhodz_new + + return tracer_new + + +@program +def vert_adv_stencil_01( + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _vert_adv_stencil_01( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + p_dtime, + out=tracer_new, + ) diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/advection/tests/test_face_val_ppm_stencil_05.py new file mode 100644 index 0000000000..c36690a5cf --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_05.py @@ -0,0 +1,83 @@ +# 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 + +import numpy as np + +from icon4py.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def face_val_ppm_stencil_05_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + z_slope: np.array, +): + p_cellhgt_mc_now_k_minus_1 = np.roll(p_cellhgt_mc_now, shift=1, axis=1) + p_cellhgt_mc_now_k_minus_2 = np.roll(p_cellhgt_mc_now, shift=2, axis=1) + p_cellhgt_mc_now_k_plus_1 = np.roll(p_cellhgt_mc_now, shift=-1, axis=1) + + p_cc_k_minus_1 = np.roll(p_cc, shift=1, axis=1) + z_slope_k_minus_1 = np.roll(z_slope, shift=1, axis=1) + + zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now_k_minus_2 + + p_cellhgt_mc_now_k_minus_1 + + p_cellhgt_mc_now + + p_cellhgt_mc_now_k_plus_1 + ) + zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( + 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( + 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 + ) + + p_face = ( + p_cc_k_minus_1 + + zgeo1 * (p_cc - p_cc_k_minus_1) + + zgeo2 + * ( + (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) + - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 + ) + ) + + return p_face + + +def test_face_val_ppm_stencil_05(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + z_slope = random_field(mesh, CellDim, KDim) + p_face = zero_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_05_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + np.asarray(z_slope), + ) + + face_val_ppm_stencil_05( + p_cc, + p_cellhgt_mc_now, + z_slope, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, 1:-1], p_face.__array__()[:, 1:-1]) diff --git a/advection/tests/test_hor_adv_stencil_01.py b/advection/tests/test_hor_adv_stencil_01.py new file mode 100644 index 0000000000..359840230a --- /dev/null +++ b/advection/tests/test_hor_adv_stencil_01.py @@ -0,0 +1,80 @@ +# 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 + +import numpy as np +from functional.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.hor_adv_stencil_01 import hor_adv_stencil_01 +from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import as_1D_sparse_field, random_field + + +def hor_adv_stencil_01_numpy( + c2e: np.array, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, + p_dtime, +) -> np.array: + geofac_div = np.expand_dims(geofac_div, axis=-1) + + tracer_new = ( + tracer_now * rhodz_now + - p_dtime * deepatmo_divh * np.sum(p_mflx_tracer_h[c2e] * geofac_div, axis=1) + ) / rhodz_new + + return tracer_new + + +def test_hor_adv_stencil_01(): + mesh = SimpleMesh() + + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + deepatmo_divh = random_field(mesh, KDim) + tracer_now = random_field(mesh, CellDim, KDim) + rhodz_now = random_field(mesh, CellDim, KDim) + rhodz_new = random_field(mesh, CellDim, KDim) + geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + tracer_new = random_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = hor_adv_stencil_01_numpy( + mesh.c2e, + np.asarray(p_mflx_tracer_h), + np.asarray(deepatmo_divh), + np.asarray(tracer_now), + np.asarray(rhodz_now), + np.asarray(rhodz_new), + np.asarray(geofac_div), + p_dtime, + ) + hor_adv_stencil_01( + p_mflx_tracer_h, + deepatmo_divh, + tracer_now, + rhodz_now, + rhodz_new, + geofac_div_new, + tracer_new, + p_dtime, + offset_provider={ + "C2E": mesh.get_c2e_offset_provider(), + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + }, + ) + assert np.allclose(tracer_new, ref) diff --git a/advection/tests/test_step_advection_stencil_03.py b/advection/tests/test_step_advection_stencil_03.py new file mode 100644 index 0000000000..ade783cb04 --- /dev/null +++ b/advection/tests/test_step_advection_stencil_03.py @@ -0,0 +1,55 @@ +# 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 + +import numpy as np + +from icon4py.advection.step_advection_stencil_03 import step_advection_stencil_03 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field + + +def step_advection_stencil_03_numpy( + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, + p_dtime, +) -> np.array: + + p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + + return p_tracer_new + + +def test_step_advection_stencil_03(): + mesh = SimpleMesh() + + p_tracer_now = random_field(mesh, CellDim, KDim) + p_grf_tend_tracer = random_field(mesh, CellDim, KDim) + p_tracer_new = random_field(mesh, CellDim, KDim) + + p_dtime = np.float64(5.0) + + ref = step_advection_stencil_03_numpy( + np.asarray(p_tracer_now), + np.asarray(p_grf_tend_tracer), + p_dtime, + ) + step_advection_stencil_03( + p_tracer_now, + p_grf_tend_tracer, + p_tracer_new, + p_dtime, + offset_provider={}, + ) + assert np.allclose(p_tracer_new, ref) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/advection/tests/test_upwind_hflux_miura_stencil_01.py new file mode 100644 index 0000000000..a4fa164a72 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -0,0 +1,94 @@ +# 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 + +import numpy as np + +from icon4py.advection.upwind_hflux_miura_stencil_01 import ( + upwind_hflux_miura_stencil_01, +) +from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def upwind_hflux_miura_stencil_01_numpy( + e2c: np.array, + p_vn: np.array, + p_cc: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + z_grad_1: np.array, + z_grad_2: np.array, + p_mass_flx_e: np.array, +) -> np.array: + + p_cc_e2c = p_cc[e2c] + z_grad_1_e2c = z_grad_1[e2c] + z_grad_2_e2c = z_grad_2[e2c] + + p_out_e = np.where( + p_vn > 0.0, + ( + p_cc_e2c[:, 0] + + distv_bary_1 * z_grad_1_e2c[:, 0] + + distv_bary_2 * z_grad_2_e2c[:, 0] + ) + * p_mass_flx_e, + ( + p_cc_e2c[:, 1] + + distv_bary_1 * z_grad_1_e2c[:, 1] + + distv_bary_2 * z_grad_2_e2c[:, 1] + ) + * p_mass_flx_e, + ) + + return p_out_e + + +def test_upwind_hflux_miura_stencil_01(): + mesh = SimpleMesh() + + p_vn = random_field(mesh, EdgeDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + distv_bary_1 = random_field(mesh, EdgeDim, KDim) + distv_bary_2 = random_field(mesh, EdgeDim, KDim) + z_grad_1 = random_field(mesh, CellDim, KDim) + z_grad_2 = random_field(mesh, CellDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_stencil_01_numpy( + mesh.e2c, + np.asarray(p_vn), + np.asarray(p_cc), + np.asarray(distv_bary_1), + np.asarray(distv_bary_2), + np.asarray(z_grad_1), + np.asarray(z_grad_2), + np.asarray(p_mass_flx_e), + ) + + upwind_hflux_miura_stencil_01( + p_vn, + p_cc, + distv_bary_1, + distv_bary_2, + z_grad_1, + z_grad_2, + p_mass_flx_e, + p_out_e, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_out_e, ref) diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py new file mode 100644 index 0000000000..c8de7c3dc2 --- /dev/null +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -0,0 +1,74 @@ +# 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 + +import numpy as np + +from icon4py.advection.vert_adv_stencil_01 import vert_adv_stencil_01 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def vert_adv_stencil_01_numpy( + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, + p_dtime, +) -> np.array: + p_mflx_tracer_v_offset_1 = np.roll(p_mflx_tracer_v, shift=-1, axis=1) + + tracer_new = ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v_offset_1 * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + ) / rhodz_new + + return tracer_new + + +def test_vert_adv_stencil_01(): + mesh = SimpleMesh() + + tracer_now = random_field(mesh, CellDim, KDim) + rhodz_now = random_field(mesh, CellDim, KDim) + p_mflx_tracer_v = random_field(mesh, CellDim, KDim) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + rhodz_new = random_field(mesh, CellDim, KDim) + tracer_new = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = vert_adv_stencil_01_numpy( + np.asarray(tracer_now), + np.asarray(rhodz_now), + np.asarray(p_mflx_tracer_v), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + np.asarray(rhodz_new), + p_dtime, + ) + vert_adv_stencil_01( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + tracer_new, + p_dtime, + offset_provider={"Koff": KDim}, + ) + assert np.allclose(tracer_new[:, :-1], ref[:, :-1]) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py new file mode 100644 index 0000000000..7367d96212 --- /dev/null +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -0,0 +1,71 @@ +# 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 + +import numpy as np + +from icon4py.advection.v_limit_prbl_sm_stencil_01 import ( + v_limit_prbl_sm_stencil_01, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def v_limit_prbl_sm_stencil_01_numpy( + p_face: np.array, + p_cc: np.array, +): + p_face_k_plus_1 = np.roll(p_face, shift=-1, axis=1) + + z_delta = p_face - p_face_k_plus_1 + z_a6i = 6.0 * (p_cc - 0.5 * (p_face + p_face_k_plus_1)) + + q_face_up, q_face_low = np.where( + abs(z_delta) < -1 * z_a6i, + np.where( + (p_cc < np.minimum(p_face, p_face_k_plus_1)), + (p_cc, p_cc), + np.where( + p_face > p_face_k_plus_1, + (3.0 * p_cc - 2.0 * p_face_k_plus_1, p_face_k_plus_1), + (p_face, 3.0 * p_cc - 2.0 * p_face), + ), + ), + (p_face, p_face_k_plus_1), + ) + + return q_face_up, q_face_low + + +def test_v_limit_prbl_sm_stencil_01(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + p_face_up = zero_field(mesh, CellDim, KDim) + p_face_low = zero_field(mesh, CellDim, KDim) + + p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_01_numpy( + np.asarray(p_face), + np.asarray(p_cc), + ) + + v_limit_prbl_sm_stencil_01( + p_face, + p_cc, + p_face_up, + p_face_low, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(p_face_up_ref[:, :-1], p_face_up[:, :-1]) + assert np.allclose(p_face_low_ref[:, :-1], p_face_low[:, :-1]) From fd9d93f08ecc67ca973a8e0ee915f9ce16714932 Mon Sep 17 00:00:00 2001 From: abishekg7 <56273301+abishekg7@users.noreply.github.com> Date: Fri, 14 Oct 2022 13:49:48 +0200 Subject: [PATCH 007/105] Fix v_limit_prbl_sm_stencil_01 (#107) * fixing return type * removing broadcast --- .../advection/v_limit_prbl_sm_stencil_01.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py index 5d06900ec9..54a9eb03e2 100644 --- a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py @@ -12,14 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import ( - Field, - FieldOffset, - abs, - broadcast, - minimum, - where, -) +from functional.ffront.fbuiltins import Field, FieldOffset, abs, minimum, where from icon4py.common.dimension import CellDim, KDim @@ -31,14 +24,11 @@ def _v_limit_prbl_sm_stencil_01( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], -) -> Field[[CellDim, KDim], float]: +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_delta = p_face - p_face(Koff[1]) z_a6i = -6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) - q_face_up = broadcast(1.0, (CellDim, KDim)) - q_face_low = broadcast(1.0, (CellDim, KDim)) - q_face_up, q_face_low = where( abs(z_delta) < z_a6i, where( From fdf3aedb1350bd94882c569d110c85e585759e55 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 19 Oct 2022 16:09:16 +0200 Subject: [PATCH 008/105] Tracer advection stencils zero copy (#106) (feat) - stencil to set field to zero for Field[[CellDim], float] and Field[[CellDim, KDim], float] --- advection/src/icon4py/advection/set_zero_c.py | 28 +++++++++++++++++++ .../src/icon4py/advection/set_zero_c_k.py | 28 +++++++++++++++++++ advection/tests/test_set_zero_c.py | 27 ++++++++++++++++++ advection/tests/test_set_zero_c_k.py | 27 ++++++++++++++++++ 4 files changed, 110 insertions(+) create mode 100644 advection/src/icon4py/advection/set_zero_c.py create mode 100644 advection/src/icon4py/advection/set_zero_c_k.py create mode 100644 advection/tests/test_set_zero_c.py create mode 100644 advection/tests/test_set_zero_c_k.py diff --git a/advection/src/icon4py/advection/set_zero_c.py b/advection/src/icon4py/advection/set_zero_c.py new file mode 100644 index 0000000000..8d8ef32c97 --- /dev/null +++ b/advection/src/icon4py/advection/set_zero_c.py @@ -0,0 +1,28 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import broadcast + +from icon4py.common.dimension import CellDim + + +@field_operator +def _set_zero_c() -> Field[[CellDim], float]: + return broadcast(0.0, (CellDim, )) + + +@program +def set_zero_c(field: Field[[CellDim], float]): + _set_zero_c(out=field) diff --git a/advection/src/icon4py/advection/set_zero_c_k.py b/advection/src/icon4py/advection/set_zero_c_k.py new file mode 100644 index 0000000000..2f30957b5c --- /dev/null +++ b/advection/src/icon4py/advection/set_zero_c_k.py @@ -0,0 +1,28 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import broadcast + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _set_zero_c_k() -> Field[[CellDim, KDim], float]: + return broadcast(0.0, (CellDim, KDim)) + + +@program +def set_zero_c_k(field: Field[[CellDim, KDim], float]): + _set_zero_c_k(out=field) diff --git a/advection/tests/test_set_zero_c.py b/advection/tests/test_set_zero_c.py new file mode 100644 index 0000000000..018cc1fc3a --- /dev/null +++ b/advection/tests/test_set_zero_c.py @@ -0,0 +1,27 @@ +# 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 + +import numpy as np + +from icon4py.advection.set_zero_c import set_zero_c +from icon4py.common.dimension import CellDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def test_set_zero_cell_k(): + mesh = SimpleMesh() + field = random_field(mesh, CellDim) + + set_zero_c(field, offset_provider={}) + assert np.allclose(field, zero_field(mesh, CellDim)) diff --git a/advection/tests/test_set_zero_c_k.py b/advection/tests/test_set_zero_c_k.py new file mode 100644 index 0000000000..1323924730 --- /dev/null +++ b/advection/tests/test_set_zero_c_k.py @@ -0,0 +1,27 @@ +# 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 + +import numpy as np + +from icon4py.advection.set_zero_c_k import set_zero_c_k +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def test_set_zero_c_k(): + mesh = SimpleMesh() + field = random_field(mesh, CellDim, KDim) + + set_zero_c_k(field, offset_provider={}) + assert np.allclose(field, zero_field(mesh, CellDim, KDim)) From 60391f193a6a39d21eb8440c9efc8aeab0c6fa15 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Oct 2022 16:15:31 +0200 Subject: [PATCH 009/105] fix pre-commit hook --- advection/src/icon4py/advection/set_zero_c.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/advection/src/icon4py/advection/set_zero_c.py b/advection/src/icon4py/advection/set_zero_c.py index 8d8ef32c97..cb9b182ea9 100644 --- a/advection/src/icon4py/advection/set_zero_c.py +++ b/advection/src/icon4py/advection/set_zero_c.py @@ -20,7 +20,7 @@ @field_operator def _set_zero_c() -> Field[[CellDim], float]: - return broadcast(0.0, (CellDim, )) + return broadcast(0.0, (CellDim,)) @program From f9df7b4a60732b7055424dd4c540d3ecdbebe53e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 25 Oct 2022 01:47:18 +0200 Subject: [PATCH 010/105] Add third coefficient to upwind hflux miura stencil 02 --- .../advection/upwind_hflux_miura_stencil_02.py | 12 +++++++----- .../tests/test_upwind_hflux_miura_stencil_02.py | 13 +++++++++---- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py index f1c8abca13..14e0739105 100644 --- a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py @@ -22,18 +22,19 @@ def _upwind_hflux_miura_stencil_02( p_cc: Field[[CellDim, KDim], float], lsq_pseudoinv_1: Field[[CECDim], float], lsq_pseudoinv_2: Field[[CECDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - p_coeff_1 = ( +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + p_coeff_1 = p_cc + p_coeff_2 = ( lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) ) - p_coeff_2 = ( + p_coeff_3 = ( lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) ) - return p_coeff_1, p_coeff_2 + return p_coeff_1, p_coeff_2, p_coeff_3 @program @@ -43,7 +44,8 @@ def upwind_hflux_miura_stencil_02( lsq_pseudoinv_2: Field[[CECDim], float], p_coeff_1: Field[[CellDim, KDim], float], p_coeff_2: Field[[CellDim, KDim], float], + p_coeff_3: Field[[CellDim, KDim], float], ): _upwind_hflux_miura_stencil_02( - p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2) + p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2, p_coeff_3) ) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py index 3dae540683..da918452e7 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_02.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -32,9 +32,10 @@ def upwind_hflux_miura_stencil_02_numpy( n_diff = p_cc[c2e2c] - p_cc_e lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - p_coeff_1 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) - p_coeff_2 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) - return p_coeff_1, p_coeff_2 + p_coeff_1 = p_cc + p_coeff_2 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_3 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return p_coeff_1, p_coeff_2, p_coeff_3 def test_upwind_hflux_miura_stencil_02(): @@ -47,8 +48,9 @@ def test_upwind_hflux_miura_stencil_02(): lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECDim) p_coeff_1 = zero_field(mesh, CellDim, KDim) p_coeff_2 = zero_field(mesh, CellDim, KDim) + p_coeff_3 = zero_field(mesh, CellDim, KDim) - ref_1, ref_2 = upwind_hflux_miura_stencil_02_numpy( + ref_1, ref_2, ref_3 = upwind_hflux_miura_stencil_02_numpy( mesh.c2e2c, np.asarray(p_cc), np.asarray(lsq_pseudoinv_1), @@ -61,6 +63,7 @@ def test_upwind_hflux_miura_stencil_02(): lsq_pseudoinv_2_field, p_coeff_1, p_coeff_2, + p_coeff_3, offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), @@ -68,5 +71,7 @@ def test_upwind_hflux_miura_stencil_02(): ) co1 = np.asarray(p_coeff_1) co2 = np.asarray(p_coeff_2) + co3 = np.asarray(p_coeff_3) assert np.allclose(ref_1, co1) assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) From b90bb4e2942a967eb649d1016d69b09b0fd01b31 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 4 Jan 2023 19:50:13 +0100 Subject: [PATCH 011/105] hflux_ffsl_hybrid (scheme 5) stencils --- .../advection/btraj_dreg_stencil_01.py | 40 ++ .../advection/btraj_dreg_stencil_02.py | 50 +++ .../advection/btraj_dreg_stencil_03.py | 146 +++++++ .../divide_flux_area_list_stencil_01.py | 370 ++++++++++++++++++ .../divide_flux_area_list_stencil_02.py | 135 +++++++ .../advection/hflux_ffsl_hybrid_stencil_02.py | 42 ++ .../prep_gauss_quadrature_c_list_stencil.py | 204 ++++++++++ .../prep_gauss_quadrature_c_stencil.py | 192 +++++++++ .../advection/recon_lsq_cell_c_svd_stencil.py | 91 +++++ 9 files changed, 1270 insertions(+) create mode 100644 advection/src/icon4py/advection/btraj_dreg_stencil_01.py create mode 100644 advection/src/icon4py/advection/btraj_dreg_stencil_02.py create mode 100644 advection/src/icon4py/advection/btraj_dreg_stencil_03.py create mode 100644 advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py create mode 100644 advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py create mode 100644 advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py create mode 100644 advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py create mode 100644 advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py create mode 100644 advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_01.py b/advection/src/icon4py/advection/btraj_dreg_stencil_01.py new file mode 100644 index 0000000000..f12200dc15 --- /dev/null +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_01.py @@ -0,0 +1,40 @@ +# 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 functional.ffront.fbuiltins import Field, int32, where, broadcast +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import EdgeDim, KDim +@field_operator +def _btraj_dreg_stencil_01( + lcounterclock: bool, + p_vn: Field[[EdgeDim, KDim], float], + tangent_orientation: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], bool]: + + tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim)) + lvn_sys_pos_true = where(p_vn*tangent_orientation >= 0.0, True, False) + mask_lcounterclock = broadcast(lcounterclock, (EdgeDim, KDim)) + lvn_sys_pos = where(mask_lcounterclock, lvn_sys_pos_true, False) + return lvn_sys_pos + + +@program +def btraj_dreg_stencil_01( + lcounterclock: bool, + p_vn: Field[[EdgeDim, KDim], float], + tangent_orientation: Field[[EdgeDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], +): + _btraj_dreg_stencil_01(lcounterclock, p_vn, tangent_orientation, out=lvn_sys_pos) + diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py new file mode 100644 index 0000000000..8aaac016ce --- /dev/null +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py @@ -0,0 +1,50 @@ +# 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 functional.ffront.fbuiltins import Field, where, sqrt, broadcast, int32 +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +@field_operator +def _btraj_dreg_stencil_02( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + edge_cell_length: Field[[ECDim], float], + p_dt: float, +) -> Field[[EdgeDim, KDim], int32]: + + lvn_pos = where(p_vn >= 0.0, True, False) + traj_length = sqrt(p_vn**2 + p_vt**2) * p_dt + e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) + famask = where(traj_length > 1.25*broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0)) + + return famask + + +@program +def btraj_dreg_stencil_02( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + edge_cell_length: Field[[ECDim], float], + famask: Field[[EdgeDim, KDim], int32], + p_dt: float, + +): + _btraj_dreg_stencil_02( + p_vn, + p_vt, + edge_cell_length, + p_dt, + out=famask) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py new file mode 100644 index 0000000000..af3fda9902 --- /dev/null +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py @@ -0,0 +1,146 @@ +# 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 functional.ffront.fbuiltins import Field, int32, where, broadcast +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +@field_operator +def _btraj_dreg_stencil_03( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + cell_idx: Field[[ECDim], int32], + cell_blk: Field[[ECDim], int32], + edge_verts_1_x: Field[[EdgeDim], float], + edge_verts_2_x: Field[[EdgeDim], float], + edge_verts_1_y: Field[[EdgeDim], float], + edge_verts_2_y: Field[[EdgeDim], float], + pos_on_tplane_e_1_x: Field[[EdgeDim], float], + pos_on_tplane_e_2_x: Field[[EdgeDim], float], + pos_on_tplane_e_1_y: Field[[EdgeDim], float], + pos_on_tplane_e_2_y: Field[[EdgeDim], float], + primal_normal_cell_x: Field[[ECDim], float], + primal_normal_cell_y: Field[[ECDim], float], + dual_normal_cell_x: Field[[ECDim], float], + dual_normal_cell_y: Field[[ECDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], + p_dt: float, +) -> tuple[ + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: +# lvn_sys_pos = where(lvn_sys_pos_int == int32(1), broadcast(True, (EdgeDim, KDim)), broadcast(False, (EdgeDim, KDim))) +# lvn_sys_pos = where(lvn_sys_pos_int == int32(1), True, False) + + lvn_pos = where(p_vn >= 0.0, True, False) + + p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) + + depart_pts_1_x = edge_verts_1_x - p_vn * p_dt + depart_pts_1_y = edge_verts_1_y - p_vt * p_dt + depart_pts_2_x = edge_verts_2_x - p_vn * p_dt + depart_pts_2_y = edge_verts_2_y - p_vt * p_dt + + pos_on_tplane_e_x = where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + pos_dreg_vert_c_2_y = where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + + pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) + pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) + dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) + dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) + + p_coords_dreg_v_1_lon = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + p_coords_dreg_v_2_lon = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + p_coords_dreg_v_3_lon = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + p_coords_dreg_v_4_lon = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + p_coords_dreg_v_1_lat = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + p_coords_dreg_v_2_lat = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + p_coords_dreg_v_3_lat = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + p_coords_dreg_v_4_lat = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + + + return p_cell_idx, p_cell_blk, p_coords_dreg_v_1_lon, p_coords_dreg_v_2_lon, p_coords_dreg_v_3_lon, p_coords_dreg_v_4_lon, p_coords_dreg_v_1_lat, p_coords_dreg_v_2_lat, p_coords_dreg_v_3_lat, p_coords_dreg_v_4_lat + + +@program +def btraj_dreg_stencil_03( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + cell_idx: Field[[ECDim], int32], + cell_blk: Field[[ECDim], int32], + edge_verts_1_x: Field[[EdgeDim], float], + edge_verts_2_x: Field[[EdgeDim], float], + edge_verts_1_y: Field[[EdgeDim], float], + edge_verts_2_y: Field[[EdgeDim], float], + pos_on_tplane_e_1_x: Field[[EdgeDim], float], + pos_on_tplane_e_2_x: Field[[EdgeDim], float], + pos_on_tplane_e_1_y: Field[[EdgeDim], float], + pos_on_tplane_e_2_y: Field[[EdgeDim], float], + primal_normal_cell_x: Field[[ECDim], float], + primal_normal_cell_y: Field[[ECDim], float], + dual_normal_cell_x: Field[[ECDim], float], + dual_normal_cell_y: Field[[ECDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], + p_dt: float, + p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_blk: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_lon: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lon: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lon: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lon: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_lat: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lat: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lat: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lat: Field[[EdgeDim, KDim], float], + +): + _btraj_dreg_stencil_03( + p_vn, + p_vt, + cell_idx, + cell_blk, + edge_verts_1_x, + edge_verts_2_x, + edge_verts_1_y, + edge_verts_2_y, + pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y, + primal_normal_cell_x, + primal_normal_cell_y, + dual_normal_cell_x, + dual_normal_cell_y, + lvn_sys_pos, + p_dt, + out=(p_cell_idx, p_cell_blk, p_coords_dreg_v_1_lon, p_coords_dreg_v_2_lon, p_coords_dreg_v_3_lon, p_coords_dreg_v_4_lon, p_coords_dreg_v_1_lat, p_coords_dreg_v_2_lat, p_coords_dreg_v_3_lat, p_coords_dreg_v_4_lat)) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py new file mode 100644 index 0000000000..bdd4480e2f --- /dev/null +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py @@ -0,0 +1,370 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, broadcast, int32, abs +from icon4py.common.dimension import E2EC, ECDim, CellDim, EdgeDim, KDim +import sys +sys.setrecursionlimit(5500) + +# FUNCTIONS +# Checking turn when travelling along three points, used to check whether lines inters. +@field_operator +def ccw( + p0_lon: Field[[EdgeDim, KDim], float], + p0_lat: Field[[EdgeDim, KDim], float], + p1_lon: Field[[EdgeDim, KDim], float], + p1_lat: Field[[EdgeDim, KDim], float], + p2_lon: Field[[EdgeDim, KDim], float], + p2_lat: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], int]: + + dx1 = p1_lon - p0_lon + dy1 = p1_lat - p0_lat + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = where(dx1dy2 > dy1dx2, True, False) + ccw_out = where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + return ccw_out + +# Checks whether two lines intersect +@field_operator +def lintersect( + line1_p1_lon: Field[[EdgeDim, KDim], float], + line1_p1_lat: Field[[EdgeDim, KDim], float], + line1_p2_lon: Field[[EdgeDim, KDim], float], + line1_p2_lat: Field[[EdgeDim, KDim], float], + line2_p1_lon: Field[[EdgeDim, KDim], float], + line2_p1_lat: Field[[EdgeDim, KDim], float], + line2_p2_lon: Field[[EdgeDim, KDim], float], + line2_p2_lat: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], bool]: + + intersect1 = ccw(line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p1_lon, line2_p1_lat) * ccw(line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p2_lon, line2_p2_lat) + intersect2 = ccw(line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, line1_p1_lon, line1_p1_lat) * ccw(line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, line1_p2_lon, line1_p2_lat) + lintersect_out = where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + +# Compute intersection point of two lines in 2D +@field_operator +def line_intersect( + line1_p1_lon: Field[[EdgeDim, KDim], float], + line1_p1_lat: Field[[EdgeDim, KDim], float], + line1_p2_lon: Field[[EdgeDim, KDim], float], + line1_p2_lat: Field[[EdgeDim, KDim], float], + line2_p1_lon: Field[[EdgeDim, KDim], float], + line2_p1_lat: Field[[EdgeDim, KDim], float], + line2_p2_lon: Field[[EdgeDim, KDim], float], + line2_p2_lat: Field[[EdgeDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + + m1 = (line1_p2_lat - line1_p1_lat)/(line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat)/(line2_p2_lon - line2_p1_lon) + + intersect_1 = (line2_p1_lat - line1_p1_lat + m1*line1_p1_lon - m2*line2_p1_lon) / (m1 - m2) + intersect_2 = line1_p1_lat + m1*(intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + +@field_operator +def _divide_flux_area_list_stencil_01( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + ptr_v3_lon: Field[[ECDim], float], + ptr_v3_lat: Field[[ECDim], float], + tangent_orientation_dsl: Field[[EdgeDim], float], + dreg_patch0_1_x: Field[[EdgeDim, KDim], float], + dreg_patch0_1_y: Field[[EdgeDim, KDim], float], + dreg_patch0_2_x: Field[[EdgeDim, KDim], float], + dreg_patch0_2_y: Field[[EdgeDim, KDim], float], + dreg_patch0_3_x: Field[[EdgeDim, KDim], float], + dreg_patch0_3_y: Field[[EdgeDim, KDim], float], + dreg_patch0_4_x: Field[[EdgeDim, KDim], float], + dreg_patch0_4_y: Field[[EdgeDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + + arrival_pts_1_x = dreg_patch0_1_x + arrival_pts_1_y = dreg_patch0_1_y + arrival_pts_2_x = dreg_patch0_2_x + arrival_pts_2_y = dreg_patch0_2_y + depart_pts_1_x = dreg_patch0_4_x # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_y = dreg_patch0_4_y + depart_pts_2_x = dreg_patch0_3_x + depart_pts_2_y = dreg_patch0_3_y + + lvn_pos = where(p_vn>=0.0, True, False) + + # get flux area departure-line segment + fl_line_p1_lon = depart_pts_1_x + fl_line_p1_lat = depart_pts_1_y + fl_line_p2_lon = depart_pts_2_x + fl_line_p2_lat = depart_pts_2_y + + # get triangle edge 1 (A1V3) + tri_line1_p1_lon = arrival_pts_1_x + tri_line1_p1_lat = arrival_pts_1_y + tri_line1_p2_lon = where(lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim))) + tri_line1_p2_lat = where(lvn_pos, broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim))) + + # get triangle edge 2 (A2V3) + tri_line2_p1_lon = arrival_pts_2_x + tri_line2_p1_lat = arrival_pts_2_y + tri_line2_p2_lon = where(lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim))) + tri_line2_p2_lat = where(lvn_pos, broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim))) + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = lintersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = lintersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + + lvn_sys_pos = where((p_vn * broadcast(tangent_orientation_dsl, (EdgeDim, KDim))) >= 0.0, True, False) + famask_bool = where(famask_int == int32(1), True, False) + #------------------------------------------------- Case 1 + mask_case1 = lintersect_line1 & lintersect_line2 & famask_bool + ps1_x, ps1_y = line_intersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + ps2_x, ps2_y = line_intersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + + # Case 1 - patch 0 + dreg_patch0_1_x = where(mask_case1, arrival_pts_1_x, dreg_patch0_1_x) + dreg_patch0_1_y = where(mask_case1, arrival_pts_1_y, dreg_patch0_1_y) + dreg_patch0_2_x = where(mask_case1, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x) + dreg_patch0_2_y = where(mask_case1, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y) + dreg_patch0_3_x = where(mask_case1, ps2_x, dreg_patch0_3_x) + dreg_patch0_3_y = where(mask_case1, ps2_y, dreg_patch0_3_y) + dreg_patch0_4_x = where(mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x) + dreg_patch0_4_y = where(mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y) + # Case 1 - patch 1 + dreg_patch1_1_x = where(mask_case1, arrival_pts_1_x, 0.0) + dreg_patch1_1_y = where(mask_case1, arrival_pts_1_y, 0.0) + dreg_patch1_4_x = where(mask_case1, arrival_pts_1_x, 0.0) + dreg_patch1_4_y = where(mask_case1, arrival_pts_1_y, 0.0) + dreg_patch1_2_x = where(mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_x), 0.0) + dreg_patch1_2_y = where(mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_y), 0.0) + dreg_patch1_3_x = where(mask_case1, where(lvn_sys_pos, depart_pts_1_x, ps1_x), 0.0) + dreg_patch1_3_y = where(mask_case1, where(lvn_sys_pos, depart_pts_1_y, ps1_y), 0.0) + # Case 1 - patch 2 + dreg_patch2_1_x = where(mask_case1, arrival_pts_2_x, 0.0) + dreg_patch2_1_y = where(mask_case1, arrival_pts_2_y, 0.0) + dreg_patch2_4_x = where(mask_case1, arrival_pts_2_x, 0.0) + dreg_patch2_4_y = where(mask_case1, arrival_pts_2_y, 0.0) + dreg_patch2_2_x = where(mask_case1, where(lvn_sys_pos, depart_pts_2_x, ps2_x), 0.0) + dreg_patch2_2_y = where(mask_case1, where(lvn_sys_pos, depart_pts_2_y, ps2_y), 0.0) + dreg_patch2_3_x = where(mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_x), 0.0) + dreg_patch2_3_y = where(mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_y), 0.0) + + #------------------------------------------------- Case 2a + mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool + # Case 2a - patch 0 + dreg_patch0_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch0_1_x) + dreg_patch0_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch0_1_y) + dreg_patch0_2_x = where(mask_case2a, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x) + dreg_patch0_2_y = where(mask_case2a, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y) + dreg_patch0_3_x = where(mask_case2a, depart_pts_2_x, dreg_patch0_3_x) + dreg_patch0_3_y = where(mask_case2a, depart_pts_2_y, dreg_patch0_3_y) + dreg_patch0_4_x = where(mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x) + dreg_patch0_4_y = where(mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y) + # Case 2a - patch 1 + dreg_patch1_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_1_x) + dreg_patch1_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_1_y) + dreg_patch1_4_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_4_x) + dreg_patch1_4_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_4_y) + dreg_patch1_2_x = where(mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_x), dreg_patch1_2_x) + dreg_patch1_2_y = where(mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_y), dreg_patch1_2_y) + dreg_patch1_3_x = where(mask_case2a, where(lvn_sys_pos, depart_pts_1_x, ps1_x), dreg_patch1_3_x) + dreg_patch1_3_y = where(mask_case2a, where(lvn_sys_pos, depart_pts_1_y, ps1_y), dreg_patch1_3_y) + # Case 2a - patch 2 + dreg_patch2_1_x = where(mask_case2a, 0.0, dreg_patch2_1_x) + dreg_patch2_1_y = where(mask_case2a, 0.0, dreg_patch2_1_y) + dreg_patch2_2_x = where(mask_case2a, 0.0, dreg_patch2_2_x) + dreg_patch2_2_y = where(mask_case2a, 0.0, dreg_patch2_2_y) + dreg_patch2_3_x = where(mask_case2a, 0.0, dreg_patch2_3_x) + dreg_patch2_3_y = where(mask_case2a, 0.0, dreg_patch2_3_y) + dreg_patch2_4_x = where(mask_case2a, 0.0, dreg_patch2_4_x) + dreg_patch2_4_y = where(mask_case2a, 0.0, dreg_patch2_4_y) + + #-------------------------------------------------- Case 2b + mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool + # Case 2b - patch 0 + dreg_patch0_1_x = where(mask_case2b, arrival_pts_1_x, dreg_patch0_1_x) + dreg_patch0_1_y = where(mask_case2b, arrival_pts_1_y, dreg_patch0_1_y) + dreg_patch0_2_x = where(mask_case2b, where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), dreg_patch0_2_x) + dreg_patch0_2_y = where(mask_case2b, where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), dreg_patch0_2_y) + dreg_patch0_3_x = where(mask_case2b, ps2_x, dreg_patch0_3_x) + dreg_patch0_3_y = where(mask_case2b, ps2_y, dreg_patch0_3_y) + dreg_patch0_4_x = where(mask_case2b, where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), dreg_patch0_4_x) + dreg_patch0_4_y = where(mask_case2b, where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), dreg_patch0_4_y) + # Case 2b - patch 1 + dreg_patch1_1_x = where(mask_case2b, 0.0, dreg_patch1_1_x) + dreg_patch1_1_y = where(mask_case2b, 0.0, dreg_patch1_1_y) + dreg_patch1_2_x = where(mask_case2b, 0.0, dreg_patch1_2_x) + dreg_patch1_2_y = where(mask_case2b, 0.0, dreg_patch1_2_y) + dreg_patch1_3_x = where(mask_case2b, 0.0, dreg_patch1_3_x) + dreg_patch1_3_y = where(mask_case2b, 0.0, dreg_patch1_3_y) + dreg_patch1_4_x = where(mask_case2b, 0.0, dreg_patch1_4_x) + dreg_patch1_4_y = where(mask_case2b, 0.0, dreg_patch1_4_y) + # Case 2b - patch 2 + dreg_patch2_1_x = where(mask_case2b, arrival_pts_2_x, dreg_patch2_1_x) + dreg_patch2_1_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_1_y) + dreg_patch2_4_x = where(mask_case2b, arrival_pts_2_x, dreg_patch2_4_x) + dreg_patch2_4_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_4_y) + dreg_patch2_2_x = where(mask_case2b, where(lvn_sys_pos, depart_pts_2_x, ps2_x), dreg_patch2_2_x) + dreg_patch2_2_y = where(mask_case2b, where(lvn_sys_pos, depart_pts_2_y, ps2_y), dreg_patch2_2_y) + dreg_patch2_3_x = where(mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_x), dreg_patch2_3_x) + dreg_patch2_3_y = where(mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_y), dreg_patch2_3_y) + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_x + fl_e1_p1_lat = arrival_pts_1_y + fl_e1_p2_lon = depart_pts_1_x + fl_e1_p2_lat = depart_pts_1_y + fl_e2_p1_lon = arrival_pts_2_x + fl_e2_p1_lat = arrival_pts_2_y + fl_e2_p2_lon = depart_pts_2_x + fl_e2_p2_lat = depart_pts_2_y + + #----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = lintersect(fl_e2_p1_lon, fl_e2_p1_lat, fl_e2_p2_lon, fl_e2_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + mask_case3a = lintersect_e2_line1 & famask_bool + pi1_x, pi1_y = line_intersect(fl_e2_p1_lon, fl_e2_p1_lat, fl_e2_p2_lon, fl_e2_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + # Case 3a - patch 0 + dreg_patch0_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch0_1_x) + dreg_patch0_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch0_1_y) + dreg_patch0_2_x = where(mask_case3a, where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), dreg_patch0_2_x) + dreg_patch0_2_y = where(mask_case3a, where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), dreg_patch0_2_y) + dreg_patch0_3_x = where(mask_case3a, ps2_x, dreg_patch0_3_x) + dreg_patch0_3_y = where(mask_case3a, ps2_y, dreg_patch0_3_y) + dreg_patch0_4_x = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), dreg_patch0_4_x) + dreg_patch0_4_y = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), dreg_patch0_4_y) + # Case 3a - patch 1 + dreg_patch1_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch1_1_x) + dreg_patch1_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch1_1_y) + dreg_patch1_2_x = where(mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_x), dreg_patch1_2_x) + dreg_patch1_2_y = where(mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_y), dreg_patch1_2_y) + dreg_patch1_3_x = where(mask_case3a, depart_pts_1_x, dreg_patch1_3_x) + dreg_patch1_3_y = where(mask_case3a, depart_pts_1_y, dreg_patch1_3_y) + dreg_patch1_4_x = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_x, pi1_x), dreg_patch1_4_x) + dreg_patch1_4_y = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_y, pi1_y), dreg_patch1_4_y) + # Case 3a - patch 2 + dreg_patch2_1_x = where(mask_case3a, 0.0, dreg_patch2_1_x) + dreg_patch2_1_y = where(mask_case3a, 0.0, dreg_patch2_1_y) + dreg_patch2_2_x = where(mask_case3a, 0.0, dreg_patch2_2_x) + dreg_patch2_2_y = where(mask_case3a, 0.0, dreg_patch2_2_y) + dreg_patch2_3_x = where(mask_case3a, 0.0, dreg_patch2_3_x) + dreg_patch2_3_y = where(mask_case3a, 0.0, dreg_patch2_3_y) + dreg_patch2_4_x = where(mask_case3a, 0.0, dreg_patch2_4_x) + dreg_patch2_4_y = where(mask_case3a, 0.0, dreg_patch2_4_y) + + #------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = lintersect(fl_e1_p1_lon, fl_e1_p1_lat, fl_e1_p2_lon, fl_e1_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = line_intersect(fl_e1_p1_lon, fl_e1_p1_lat, fl_e1_p2_lon, fl_e1_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + # Case 3b - patch 0 + dreg_patch0_1_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_1_x) + dreg_patch0_1_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_1_y) + dreg_patch0_4_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_4_x) + dreg_patch0_4_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_4_y) + dreg_patch0_2_x = where(mask_case3b, where(lvn_sys_pos, arrival_pts_2_x, pi2_x), dreg_patch0_2_x) + dreg_patch0_2_y = where(mask_case3b, where(lvn_sys_pos, arrival_pts_2_y, pi2_y), dreg_patch0_2_y) + dreg_patch0_3_x = where(mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_x), dreg_patch0_3_x) + dreg_patch0_3_y = where(mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_y), dreg_patch0_3_y) + # Case 3b - patch 1 + dreg_patch1_1_x = where(mask_case3b, 0.0, dreg_patch1_1_x) + dreg_patch1_1_y = where(mask_case3b, 0.0, dreg_patch1_1_y) + dreg_patch1_2_x = where(mask_case3b, 0.0, dreg_patch1_2_x) + dreg_patch1_2_y = where(mask_case3b, 0.0, dreg_patch1_2_y) + dreg_patch1_3_x = where(mask_case3b, 0.0, dreg_patch1_3_x) + dreg_patch1_3_y = where(mask_case3b, 0.0, dreg_patch1_3_y) + dreg_patch1_4_x = where(mask_case3b, 0.0, dreg_patch1_4_x) + dreg_patch1_4_y = where(mask_case3b, 0.0, dreg_patch1_4_y) + # Case 3b - patch 2 + dreg_patch2_1_x = where(mask_case3b, arrival_pts_2_x, dreg_patch2_1_x) + dreg_patch2_1_y = where(mask_case3b, arrival_pts_2_y, dreg_patch2_1_y) + dreg_patch2_2_x = where(mask_case3b, where(lvn_sys_pos, depart_pts_2_x, pi2_x), dreg_patch2_2_x) + dreg_patch2_2_y = where(mask_case3b, where(lvn_sys_pos, depart_pts_2_y, pi2_y), dreg_patch2_2_y) + dreg_patch2_3_x = where(mask_case3b, depart_pts_1_x, dreg_patch2_3_x) + dreg_patch2_3_y = where(mask_case3b, depart_pts_1_y, dreg_patch2_3_y) + dreg_patch2_4_x = where(mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_x), dreg_patch2_4_x) + dreg_patch2_4_y = where(mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_y), dreg_patch2_4_y) + + #--------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = (mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1) +# mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = famask_bool & (not indices_previously_matched) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_x = where(mask_case4, 0.0, dreg_patch1_1_x) + dreg_patch1_1_y = where(mask_case4, 0.0, dreg_patch1_1_y) + dreg_patch1_2_x = where(mask_case4, 0.0, dreg_patch1_2_x) + dreg_patch1_2_y = where(mask_case4, 0.0, dreg_patch1_2_y) + dreg_patch1_3_x = where(mask_case4, 0.0, dreg_patch1_3_x) + dreg_patch1_3_y = where(mask_case4, 0.0, dreg_patch1_3_y) + dreg_patch1_4_x = where(mask_case4, 0.0, dreg_patch1_4_x) + dreg_patch1_4_y = where(mask_case4, 0.0, dreg_patch1_4_y) + # Case 4 - patch 2 + dreg_patch2_1_x = where(mask_case4, 0.0, dreg_patch2_1_x) + dreg_patch2_1_y = where(mask_case4, 0.0, dreg_patch2_1_y) + dreg_patch2_2_x = where(mask_case4, 0.0, dreg_patch2_2_x) + dreg_patch2_2_y = where(mask_case4, 0.0, dreg_patch2_2_y) + dreg_patch2_3_x = where(mask_case4, 0.0, dreg_patch2_3_x) + dreg_patch2_3_y = where(mask_case4, 0.0, dreg_patch2_3_y) + dreg_patch2_4_x = where(mask_case4, 0.0, dreg_patch2_4_x) + dreg_patch2_4_y = where(mask_case4, 0.0, dreg_patch2_4_y) + + + return dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y + +@program +def divide_flux_area_list_stencil_01( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + ptr_v3_lon: Field[[ECDim], float], + ptr_v3_lat: Field[[ECDim], float], + tangent_orientation_dsl: Field[[EdgeDim], float], + dreg_patch0_1_x: Field[[EdgeDim, KDim], float], + dreg_patch0_1_y: Field[[EdgeDim, KDim], float], + dreg_patch0_2_x: Field[[EdgeDim, KDim], float], + dreg_patch0_2_y: Field[[EdgeDim, KDim], float], + dreg_patch0_3_x: Field[[EdgeDim, KDim], float], + dreg_patch0_3_y: Field[[EdgeDim, KDim], float], + dreg_patch0_4_x: Field[[EdgeDim, KDim], float], + dreg_patch0_4_y: Field[[EdgeDim, KDim], float], + dreg_patch1_1_x: Field[[EdgeDim, KDim], float], + dreg_patch1_1_y: Field[[EdgeDim, KDim], float], + dreg_patch1_2_x: Field[[EdgeDim, KDim], float], + dreg_patch1_2_y: Field[[EdgeDim, KDim], float], + dreg_patch1_3_x: Field[[EdgeDim, KDim], float], + dreg_patch1_3_y: Field[[EdgeDim, KDim], float], + dreg_patch1_4_x: Field[[EdgeDim, KDim], float], + dreg_patch1_4_y: Field[[EdgeDim, KDim], float], + dreg_patch2_1_x: Field[[EdgeDim, KDim], float], + dreg_patch2_1_y: Field[[EdgeDim, KDim], float], + dreg_patch2_2_x: Field[[EdgeDim, KDim], float], + dreg_patch2_2_y: Field[[EdgeDim, KDim], float], + dreg_patch2_3_x: Field[[EdgeDim, KDim], float], + dreg_patch2_3_y: Field[[EdgeDim, KDim], float], + dreg_patch2_4_x: Field[[EdgeDim, KDim], float], + dreg_patch2_4_y: Field[[EdgeDim, KDim], float], +): + _divide_flux_area_list_stencil_01( + famask_int, p_vn, ptr_v3_lon, ptr_v3_lat, tangent_orientation_dsl, dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, out=(dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y) + ) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py new file mode 100644 index 0000000000..aae927a481 --- /dev/null +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py @@ -0,0 +1,135 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, broadcast, int32, abs +from icon4py.common.dimension import E2EC, ECDim, CellDim, EdgeDim, KDim +import sys +sys.setrecursionlimit(5500) +@field_operator +def _divide_flux_area_list_stencil_02( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + bf_cc_patch1_lon: Field[[ECDim], float], + bf_cc_patch1_lat: Field[[ECDim], float], + bf_cc_patch2_lon: Field[[ECDim], float], + bf_cc_patch2_lat: Field[[ECDim], float], + butterfly_idx_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], + dreg_patch1_1_x: Field[[EdgeDim, KDim], float], + dreg_patch1_1_y: Field[[EdgeDim, KDim], float], + dreg_patch1_2_x: Field[[EdgeDim, KDim], float], + dreg_patch1_2_y: Field[[EdgeDim, KDim], float], + dreg_patch1_3_x: Field[[EdgeDim, KDim], float], + dreg_patch1_3_y: Field[[EdgeDim, KDim], float], + dreg_patch1_4_x: Field[[EdgeDim, KDim], float], + dreg_patch1_4_y: Field[[EdgeDim, KDim], float], + dreg_patch2_1_x: Field[[EdgeDim, KDim], float], + dreg_patch2_1_y: Field[[EdgeDim, KDim], float], + dreg_patch2_2_x: Field[[EdgeDim, KDim], float], + dreg_patch2_2_y: Field[[EdgeDim, KDim], float], + dreg_patch2_3_x: Field[[EdgeDim, KDim], float], + dreg_patch2_3_y: Field[[EdgeDim, KDim], float], + dreg_patch2_4_x: Field[[EdgeDim, KDim], float], + dreg_patch2_4_y: Field[[EdgeDim, KDim], float], +) -> tuple[ Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32]]: + famask_bool = where(famask_int == int32(1), True, False) + lvn_pos = where(p_vn >= 0.0, True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon = where(famask_bool, where(lvn_pos, bf_cc_patch1_lon(E2EC[0]), bf_cc_patch1_lon(E2EC[1])), 0.0) + bf_cc_patch1_lat = where(famask_bool, where(lvn_pos, bf_cc_patch1_lat(E2EC[0]), bf_cc_patch1_lat(E2EC[1])), 0.0) + bf_cc_patch2_lon = where(famask_bool, where(lvn_pos, bf_cc_patch2_lon(E2EC[0]), bf_cc_patch2_lon(E2EC[1])), 0.0) + bf_cc_patch2_lat = where(famask_bool, where(lvn_pos, bf_cc_patch2_lat(E2EC[0]), bf_cc_patch2_lat(E2EC[1])), 0.0) + + # patch1 in translated system + dreg_patch1_1_x = dreg_patch1_1_x - bf_cc_patch1_lon + dreg_patch1_1_y = dreg_patch1_1_y - bf_cc_patch1_lat + dreg_patch1_2_x = dreg_patch1_2_x - bf_cc_patch1_lon + dreg_patch1_2_y = dreg_patch1_2_y - bf_cc_patch1_lat + dreg_patch1_3_x = dreg_patch1_3_x - bf_cc_patch1_lon + dreg_patch1_3_y = dreg_patch1_3_y - bf_cc_patch1_lat + dreg_patch1_4_x = dreg_patch1_4_x - bf_cc_patch1_lon + dreg_patch1_4_y = dreg_patch1_4_y - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_x = dreg_patch2_1_x - bf_cc_patch2_lon + dreg_patch2_1_y = dreg_patch2_1_y - bf_cc_patch2_lat + dreg_patch2_2_x = dreg_patch2_2_x - bf_cc_patch2_lon + dreg_patch2_2_y = dreg_patch2_2_y - bf_cc_patch2_lat + dreg_patch2_3_x = dreg_patch2_3_x - bf_cc_patch2_lon + dreg_patch2_3_y = dreg_patch2_3_y - bf_cc_patch2_lat + dreg_patch2_4_x = dreg_patch2_4_x - bf_cc_patch2_lon + dreg_patch2_4_y = dreg_patch2_4_y - bf_cc_patch2_lat + + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = broadcast(butterfly_idx_patch1_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch1_vnneg_3d = broadcast(butterfly_idx_patch1_vnneg, (EdgeDim, KDim)) + butterfly_idx_patch2_vnpos_3d = broadcast(butterfly_idx_patch2_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch2_vnneg_3d = broadcast(butterfly_idx_patch2_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch1_vnpos_3d = broadcast(butterfly_blk_patch1_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch1_vnneg_3d = broadcast(butterfly_blk_patch1_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch2_vnpos_3d = broadcast(butterfly_blk_patch2_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch2_vnneg_3d = broadcast(butterfly_blk_patch2_vnneg, (EdgeDim, KDim)) + patch1_cell_idx_dsl = where(famask_bool, where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), int32(0)) + patch2_cell_idx_dsl = where(famask_bool, where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), int32(0)) + patch1_cell_blk_dsl = where(famask_bool, where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), int32(0)) + patch2_cell_blk_dsl = where(famask_bool, where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), int32(0)) + + return dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, patch1_cell_idx_dsl, patch1_cell_blk_dsl, patch2_cell_idx_dsl, patch2_cell_blk_dsl + +@program +def divide_flux_area_list_stencil_02( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + bf_cc_patch1_lon: Field[[ECDim], float], + bf_cc_patch1_lat: Field[[ECDim], float], + bf_cc_patch2_lon: Field[[ECDim], float], + bf_cc_patch2_lat: Field[[ECDim], float], + butterfly_idx_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], + dreg_patch1_1_x: Field[[EdgeDim, KDim], float], + dreg_patch1_1_y: Field[[EdgeDim, KDim], float], + dreg_patch1_2_x: Field[[EdgeDim, KDim], float], + dreg_patch1_2_y: Field[[EdgeDim, KDim], float], + dreg_patch1_3_x: Field[[EdgeDim, KDim], float], + dreg_patch1_3_y: Field[[EdgeDim, KDim], float], + dreg_patch1_4_x: Field[[EdgeDim, KDim], float], + dreg_patch1_4_y: Field[[EdgeDim, KDim], float], + dreg_patch2_1_x: Field[[EdgeDim, KDim], float], + dreg_patch2_1_y: Field[[EdgeDim, KDim], float], + dreg_patch2_2_x: Field[[EdgeDim, KDim], float], + dreg_patch2_2_y: Field[[EdgeDim, KDim], float], + dreg_patch2_3_x: Field[[EdgeDim, KDim], float], + dreg_patch2_3_y: Field[[EdgeDim, KDim], float], + dreg_patch2_4_x: Field[[EdgeDim, KDim], float], + dreg_patch2_4_y: Field[[EdgeDim, KDim], float], + patch1_cell_idx_dsl: Field[[EdgeDim, KDim], int32], + patch1_cell_blk_dsl: Field[[EdgeDim, KDim], int32], + patch2_cell_idx_dsl: Field[[EdgeDim, KDim], int32], + patch2_cell_blk_dsl: Field[[EdgeDim, KDim], int32], +): + _divide_flux_area_list_stencil_02( + famask_int, p_vn, bf_cc_patch1_lon, bf_cc_patch1_lat, bf_cc_patch2_lon, bf_cc_patch2_lat, butterfly_idx_patch1_vnpos, butterfly_idx_patch1_vnneg, butterfly_blk_patch1_vnpos, butterfly_blk_patch1_vnneg, butterfly_idx_patch2_vnpos, butterfly_idx_patch2_vnneg, butterfly_blk_patch2_vnpos, butterfly_blk_patch2_vnneg, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, out=( dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, patch1_cell_idx_dsl, patch1_cell_blk_dsl, patch2_cell_idx_dsl, patch2_cell_blk_dsl) + ) diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py new file mode 100644 index 0000000000..ba22cf6116 --- /dev/null +++ b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py @@ -0,0 +1,42 @@ +# 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 functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field +from icon4py.common.dimension import EdgeDim, KDim + + +@field_operator +def _hflux_ffsl_hybrid_stencil_02( + p_out_e: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float] +) -> Field[[EdgeDim, KDim], float]: + + p_out_e = p_mass_flx_e * p_out_e / z_dreg_area + + return p_out_e + + +@program +def hflux_ffsl_hybrid_stencil_02( + p_out_e: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float] +): + _hflux_ffsl_hybrid_stencil_02( + p_out_e, + p_mass_flx_e, + z_dreg_area, + out=p_out_e, + ) diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py new file mode 100644 index 0000000000..e15418bc04 --- /dev/null +++ b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py @@ -0,0 +1,204 @@ +# 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 functional.ffront.fbuiltins import Field, maximum, abs, where, broadcast, int32 +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import EdgeDim, KDim + + +@field_operator +def _prep_gauss_quadrature_c_list_stencil( + famask_int: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area: Field[[EdgeDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float] +]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + famask_bool = where(famask_int == int32(1), True, False) + p_coords_dreg_v_1_x = where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = where(famask_bool, p_coords_dreg_v_4_y, 0.0) + + wgt_t_detjac_1 = where(famask_bool, dbl_eps + z_wgt_1 * ( + ( z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) + wgt_t_detjac_2 = where(famask_bool, dbl_eps + z_wgt_2 * ( + ( z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) + wgt_t_detjac_3 = where(famask_bool, dbl_eps + z_wgt_3 * ( + ( z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) + wgt_t_detjac_4 = where(famask_bool, dbl_eps + z_wgt_4 * ( + ( z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) + + z_gauss_pts_1_x = shape_func_1_1 * p_coords_dreg_v_1_x + shape_func_2_1 * p_coords_dreg_v_2_x + shape_func_3_1 * p_coords_dreg_v_3_x + shape_func_4_1 * p_coords_dreg_v_4_x + z_gauss_pts_1_y = shape_func_1_1 * p_coords_dreg_v_1_y + shape_func_2_1 * p_coords_dreg_v_2_y + shape_func_3_1 * p_coords_dreg_v_3_y + shape_func_4_1 * p_coords_dreg_v_4_y + z_gauss_pts_2_x = shape_func_1_2 * p_coords_dreg_v_1_x + shape_func_2_2 * p_coords_dreg_v_2_x + shape_func_3_2 * p_coords_dreg_v_3_x + shape_func_4_2 * p_coords_dreg_v_4_x + z_gauss_pts_2_y = shape_func_1_2 * p_coords_dreg_v_1_y + shape_func_2_2 * p_coords_dreg_v_2_y + shape_func_3_2 * p_coords_dreg_v_3_y + shape_func_4_2 * p_coords_dreg_v_4_y + z_gauss_pts_3_x = shape_func_1_3 * p_coords_dreg_v_1_x + shape_func_2_3 * p_coords_dreg_v_2_x + shape_func_3_3 * p_coords_dreg_v_3_x + shape_func_4_3 * p_coords_dreg_v_4_x + z_gauss_pts_3_y = shape_func_1_3 * p_coords_dreg_v_1_y + shape_func_2_3 * p_coords_dreg_v_2_y + shape_func_3_3 * p_coords_dreg_v_3_y + shape_func_4_3 * p_coords_dreg_v_4_y + z_gauss_pts_4_x = shape_func_1_4 * p_coords_dreg_v_1_x + shape_func_2_4 * p_coords_dreg_v_2_x + shape_func_3_4 * p_coords_dreg_v_3_x + shape_func_4_4 * p_coords_dreg_v_4_x + z_gauss_pts_4_y = shape_func_1_4 * p_coords_dreg_v_1_y + shape_func_2_4 * p_coords_dreg_v_2_y + shape_func_3_4 * p_coords_dreg_v_3_y + shape_func_4_4 * p_coords_dreg_v_4_y + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x + wgt_t_detjac_3 * z_gauss_pts_3_x + wgt_t_detjac_4 * z_gauss_pts_4_x + p_quad_vector_sum_3 = wgt_t_detjac_1 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_y + p_quad_vector_sum_4 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + p_quad_vector_sum_5 = wgt_t_detjac_1 * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + p_quad_vector_sum_6 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + p_quad_vector_sum_7 = wgt_t_detjac_1 * z_gauss_pts_1_x**3 + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + p_quad_vector_sum_8 = wgt_t_detjac_1 * z_gauss_pts_1_y**3 + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + p_quad_vector_sum_9 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + p_quad_vector_sum_10 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + + p_dreg_area = p_dreg_area + p_quad_vector_sum_1 + + + return p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3,p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10 + + +@program +def prep_gauss_quadrature_c_list_stencil( + famask_int: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], +): + _prep_gauss_quadrature_c_list_stencil( + famask_int, p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, p_coords_dreg_v_4_x, p_coords_dreg_v_1_y, p_coords_dreg_v_2_y, p_coords_dreg_v_3_y, p_coords_dreg_v_4_y, shape_func_1_1, shape_func_2_1, shape_func_3_1, shape_func_4_1, shape_func_1_2, shape_func_2_2, shape_func_3_2, shape_func_4_2, shape_func_1_3, shape_func_2_3, shape_func_3_3, shape_func_4_3, shape_func_1_4, shape_func_2_4, shape_func_3_4, shape_func_4_4, zeta_1, zeta_2, zeta_3, zeta_4, eta_1, eta_2, eta_3, eta_4, wgt_zeta_1, wgt_zeta_2, wgt_eta_1, wgt_eta_2, dbl_eps, eps, p_dreg_area, out=(p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10) + ) diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py b/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py new file mode 100644 index 0000000000..9efa963c18 --- /dev/null +++ b/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py @@ -0,0 +1,192 @@ +# 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 functional.ffront.fbuiltins import Field, maximum, abs, where, broadcast +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import EdgeDim, KDim + + +@field_operator +def _prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float] +]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) + * (z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) + - (z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) + * (z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + + z_gauss_pts_1_x = shape_func_1_1 * p_coords_dreg_v_1_x + shape_func_2_1 * p_coords_dreg_v_2_x + shape_func_3_1 * p_coords_dreg_v_3_x + shape_func_4_1 * p_coords_dreg_v_4_x + z_gauss_pts_1_y = shape_func_1_1 * p_coords_dreg_v_1_y + shape_func_2_1 * p_coords_dreg_v_2_y + shape_func_3_1 * p_coords_dreg_v_3_y + shape_func_4_1 * p_coords_dreg_v_4_y + z_gauss_pts_2_x = shape_func_1_2 * p_coords_dreg_v_1_x + shape_func_2_2 * p_coords_dreg_v_2_x + shape_func_3_2 * p_coords_dreg_v_3_x + shape_func_4_2 * p_coords_dreg_v_4_x + z_gauss_pts_2_y = shape_func_1_2 * p_coords_dreg_v_1_y + shape_func_2_2 * p_coords_dreg_v_2_y + shape_func_3_2 * p_coords_dreg_v_3_y + shape_func_4_2 * p_coords_dreg_v_4_y + z_gauss_pts_3_x = shape_func_1_3 * p_coords_dreg_v_1_x + shape_func_2_3 * p_coords_dreg_v_2_x + shape_func_3_3 * p_coords_dreg_v_3_x + shape_func_4_3 * p_coords_dreg_v_4_x + z_gauss_pts_3_y = shape_func_1_3 * p_coords_dreg_v_1_y + shape_func_2_3 * p_coords_dreg_v_2_y + shape_func_3_3 * p_coords_dreg_v_3_y + shape_func_4_3 * p_coords_dreg_v_4_y + z_gauss_pts_4_x = shape_func_1_4 * p_coords_dreg_v_1_x + shape_func_2_4 * p_coords_dreg_v_2_x + shape_func_3_4 * p_coords_dreg_v_3_x + shape_func_4_4 * p_coords_dreg_v_4_x + z_gauss_pts_4_y = shape_func_1_4 * p_coords_dreg_v_1_y + shape_func_2_4 * p_coords_dreg_v_2_y + shape_func_3_4 * p_coords_dreg_v_3_y + shape_func_4_4 * p_coords_dreg_v_4_y + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x + wgt_t_detjac_3 * z_gauss_pts_3_x + wgt_t_detjac_4 * z_gauss_pts_4_x + p_quad_vector_sum_3 = wgt_t_detjac_1 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_y + p_quad_vector_sum_4 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + p_quad_vector_sum_5 = wgt_t_detjac_1 * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + p_quad_vector_sum_6 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + p_quad_vector_sum_7 = wgt_t_detjac_1 * z_gauss_pts_1_x**3 + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + p_quad_vector_sum_8 = wgt_t_detjac_1 * z_gauss_pts_1_y**3 + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + p_quad_vector_sum_9 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + p_quad_vector_sum_10 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = where(z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area))) + + + return p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3,p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6,p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9,p_quad_vector_sum_10, p_dreg_area_out + + +@program +def prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + p_dreg_area_out: Field[[EdgeDim, KDim], float], +): + _prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, p_coords_dreg_v_4_x, p_coords_dreg_v_1_y, p_coords_dreg_v_2_y, p_coords_dreg_v_3_y, p_coords_dreg_v_4_y, shape_func_1_1, shape_func_2_1, shape_func_3_1, shape_func_4_1, shape_func_1_2, shape_func_2_2, shape_func_3_2, shape_func_4_2, shape_func_1_3, shape_func_2_3, shape_func_3_3, shape_func_4_3, shape_func_1_4, shape_func_2_4, shape_func_3_4, shape_func_4_4, zeta_1, zeta_2, zeta_3, zeta_4, eta_1, eta_2, eta_3, eta_4, wgt_zeta_1, wgt_zeta_2, wgt_eta_1, wgt_eta_2, dbl_eps, eps, out=(p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, p_dreg_area_out) + ) diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py new file mode 100644 index 0000000000..e81f7550e3 --- /dev/null +++ b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py @@ -0,0 +1,91 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, C2E2C2E2C, CECECDim, CellDim, KDim + + +@field_operator +def _recon_lsq_cell_c_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECECDim], float], + lsq_pseudoinv_2: Field[[CECECDim], float], + lsq_pseudoinv_3: Field[[CECECDim], float], + lsq_pseudoinv_4: Field[[CECECDim], float], + lsq_pseudoinv_5: Field[[CECECDim], float], + lsq_pseudoinv_6: Field[[CECECDim], float], + lsq_pseudoinv_7: Field[[CECECDim], float], + lsq_pseudoinv_8: Field[[CECECDim], float], + lsq_pseudoinv_9: Field[[CECECDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + p_coeff_10 = lsq_pseudoinv_9(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_9(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_9(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_9(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_9(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_9(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_9(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_9(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_9(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_9 = lsq_pseudoinv_8(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_8(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_8(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_8(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_8(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_8(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_8(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_8(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_8(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_8 = lsq_pseudoinv_7(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_7(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_7(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_7(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_7(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_7(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_7(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_7(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_7(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_7 = lsq_pseudoinv_6(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_6(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_6(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_6(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_6(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_6(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_6(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_6(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_6(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_6 = lsq_pseudoinv_5(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_5(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_5(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_5(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_5(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_5(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_5(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_5(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_5(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_5 = lsq_pseudoinv_4(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_4(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_4(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_4(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_4(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_4(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_4(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_4(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_4(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_4 = lsq_pseudoinv_3(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_3(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_3(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_3(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_3(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_3(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_3(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_3(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_3(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_3 = lsq_pseudoinv_2(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_2(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_2(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_2(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_2(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_2(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_2(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_2(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_2(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_2 = lsq_pseudoinv_1(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_1(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_1(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_1(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_1(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_1(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_1(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_1(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_1(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + + p_coeff_1 = p_cc - (p_coeff_2*lsq_moments_1 + p_coeff_3*lsq_moments_2 + p_coeff_4*lsq_moments_3 + p_coeff_5*lsq_moments_4 + p_coeff_6*lsq_moments_5 + p_coeff_7*lsq_moments_6 + p_coeff_8*lsq_moments_7 + p_coeff_9*lsq_moments_8 + p_coeff_10*lsq_moments_9) + + return p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10 + + +@program +def recon_lsq_cell_c_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECECDim], float], + lsq_pseudoinv_2: Field[[CECECDim], float], + lsq_pseudoinv_3: Field[[CECECDim], float], + lsq_pseudoinv_4: Field[[CECECDim], float], + lsq_pseudoinv_5: Field[[CECECDim], float], + lsq_pseudoinv_6: Field[[CECECDim], float], + lsq_pseudoinv_7: Field[[CECECDim], float], + lsq_pseudoinv_8: Field[[CECECDim], float], + lsq_pseudoinv_9: Field[[CECECDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], + p_coeff_1: Field[[CellDim, KDim], float], + p_coeff_2: Field[[CellDim, KDim], float], + p_coeff_3: Field[[CellDim, KDim], float], + p_coeff_4: Field[[CellDim, KDim], float], + p_coeff_5: Field[[CellDim, KDim], float], + p_coeff_6: Field[[CellDim, KDim], float], + p_coeff_7: Field[[CellDim, KDim], float], + p_coeff_8: Field[[CellDim, KDim], float], + p_coeff_9: Field[[CellDim, KDim], float], + p_coeff_10: Field[[CellDim, KDim], float], +): + _recon_lsq_cell_c_svd_stencil( + p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, lsq_pseudoinv_3, lsq_pseudoinv_4, lsq_pseudoinv_5, lsq_pseudoinv_6, lsq_pseudoinv_7, lsq_pseudoinv_8, lsq_pseudoinv_9, lsq_moments_1, lsq_moments_2, lsq_moments_3, lsq_moments_4, lsq_moments_5, lsq_moments_6, lsq_moments_7, lsq_moments_8, lsq_moments_9, out=(p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10)) From 98c60e6b4c033083c090fdeb9a756b8235351585 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 6 Jan 2023 16:54:13 +0100 Subject: [PATCH 012/105] recon_lsq_cell_c_stencil.py addition but svd=False doesn't work in ICON --- .../advection/recon_lsq_cell_c_stencil.py | 207 ++++++++++++++++++ 1 file changed, 207 insertions(+) create mode 100644 advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py new file mode 100644 index 0000000000..4b1862672c --- /dev/null +++ b/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py @@ -0,0 +1,207 @@ +# 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 functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import broadcast +from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim + +import sys # Increase recusion depth, otherwise it doesn't compile +sys.setrecursionlimit(6000) + +@field_operator +def _recon_lsq_cell_c_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_qtmat_c_1: Field[[CECECDim], float], + lsq_qtmat_c_2: Field[[CECECDim], float], + lsq_qtmat_c_3: Field[[CECECDim], float], + lsq_qtmat_c_4: Field[[CECECDim], float], + lsq_qtmat_c_5: Field[[CECECDim], float], + lsq_qtmat_c_6: Field[[CECECDim], float], + lsq_qtmat_c_7: Field[[CECECDim], float], + lsq_qtmat_c_8: Field[[CECECDim], float], + lsq_qtmat_c_9: Field[[CECECDim], float], + lsq_rmat_rdiag_c_1: Field[[CellDim], float], + lsq_rmat_rdiag_c_2: Field[[CellDim], float], + lsq_rmat_rdiag_c_3: Field[[CellDim], float], + lsq_rmat_rdiag_c_4: Field[[CellDim], float], + lsq_rmat_rdiag_c_5: Field[[CellDim], float], + lsq_rmat_rdiag_c_6: Field[[CellDim], float], + lsq_rmat_rdiag_c_7: Field[[CellDim], float], + lsq_rmat_rdiag_c_8: Field[[CellDim], float], + lsq_rmat_rdiag_c_9: Field[[CellDim], float], + lsq_rmat_utri_c_1: Field[[CellDim], float], + lsq_rmat_utri_c_2: Field[[CellDim], float], + lsq_rmat_utri_c_3: Field[[CellDim], float], + lsq_rmat_utri_c_4: Field[[CellDim], float], + lsq_rmat_utri_c_5: Field[[CellDim], float], + lsq_rmat_utri_c_6: Field[[CellDim], float], + lsq_rmat_utri_c_7: Field[[CellDim], float], + lsq_rmat_utri_c_8: Field[[CellDim], float], + lsq_rmat_utri_c_9: Field[[CellDim], float], + lsq_rmat_utri_c_10: Field[[CellDim], float], + lsq_rmat_utri_c_11: Field[[CellDim], float], + lsq_rmat_utri_c_12: Field[[CellDim], float], + lsq_rmat_utri_c_13: Field[[CellDim], float], + lsq_rmat_utri_c_14: Field[[CellDim], float], + lsq_rmat_utri_c_15: Field[[CellDim], float], + lsq_rmat_utri_c_16: Field[[CellDim], float], + lsq_rmat_utri_c_17: Field[[CellDim], float], + lsq_rmat_utri_c_18: Field[[CellDim], float], + lsq_rmat_utri_c_19: Field[[CellDim], float], + lsq_rmat_utri_c_20: Field[[CellDim], float], + lsq_rmat_utri_c_21: Field[[CellDim], float], + lsq_rmat_utri_c_22: Field[[CellDim], float], + lsq_rmat_utri_c_23: Field[[CellDim], float], + lsq_rmat_utri_c_24: Field[[CellDim], float], + lsq_rmat_utri_c_25: Field[[CellDim], float], + lsq_rmat_utri_c_26: Field[[CellDim], float], + lsq_rmat_utri_c_27: Field[[CellDim], float], + lsq_rmat_utri_c_28: Field[[CellDim], float], + lsq_rmat_utri_c_29: Field[[CellDim], float], + lsq_rmat_utri_c_30: Field[[CellDim], float], + lsq_rmat_utri_c_31: Field[[CellDim], float], + lsq_rmat_utri_c_32: Field[[CellDim], float], + lsq_rmat_utri_c_33: Field[[CellDim], float], + lsq_rmat_utri_c_34: Field[[CellDim], float], + lsq_rmat_utri_c_35: Field[[CellDim], float], + lsq_rmat_utri_c_36: Field[[CellDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + lsq_rmat_rdiag_c_1 = broadcast(lsq_rmat_rdiag_c_1, (CellDim, KDim)) + lsq_rmat_rdiag_c_2 = broadcast(lsq_rmat_rdiag_c_2, (CellDim, KDim)) + lsq_rmat_rdiag_c_3 = broadcast(lsq_rmat_rdiag_c_3, (CellDim, KDim)) + lsq_rmat_rdiag_c_4 = broadcast(lsq_rmat_rdiag_c_4, (CellDim, KDim)) + lsq_rmat_rdiag_c_5 = broadcast(lsq_rmat_rdiag_c_5, (CellDim, KDim)) + lsq_rmat_rdiag_c_6 = broadcast(lsq_rmat_rdiag_c_6, (CellDim, KDim)) + lsq_rmat_rdiag_c_7 = broadcast(lsq_rmat_rdiag_c_7, (CellDim, KDim)) + lsq_rmat_rdiag_c_8 = broadcast(lsq_rmat_rdiag_c_8, (CellDim, KDim)) + lsq_rmat_rdiag_c_9 = broadcast(lsq_rmat_rdiag_c_9, (CellDim, KDim)) + lsq_qtmat_c_1 = broadcast(lsq_qtmat_c_1, (CECECDim, KDim)) + lsq_qtmat_c_2 = broadcast(lsq_qtmat_c_2, (CECECDim, KDim)) + lsq_qtmat_c_3 = broadcast(lsq_qtmat_c_3, (CECECDim, KDim)) + lsq_qtmat_c_4 = broadcast(lsq_qtmat_c_4, (CECECDim, KDim)) + lsq_qtmat_c_5 = broadcast(lsq_qtmat_c_5, (CECECDim, KDim)) + lsq_qtmat_c_6 = broadcast(lsq_qtmat_c_6, (CECECDim, KDim)) + lsq_qtmat_c_7 = broadcast(lsq_qtmat_c_7, (CECECDim, KDim)) + lsq_qtmat_c_8 = broadcast(lsq_qtmat_c_8, (CECECDim, KDim)) + lsq_qtmat_c_9 = broadcast(lsq_qtmat_c_9, (CECECDim, KDim)) + + p_coeff_10 = lsq_rmat_rdiag_c_9 * (lsq_qtmat_c_9(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_9(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_9(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_9(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_9(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_9(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_9(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_9(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_9(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc)) + + p_coeff_9 = lsq_rmat_rdiag_c_8 * (lsq_qtmat_c_8(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_8(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_8(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_8(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_8(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_8(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_8(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_8(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_8(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - lsq_rmat_utri_c_1*p_coeff_10) + + p_coeff_8 = lsq_rmat_rdiag_c_7 * (lsq_qtmat_c_7(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_7(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_7(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_7(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_7(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_7(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_7(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_7(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_7(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_2*p_coeff_9 + lsq_rmat_utri_c_3*p_coeff_10)) + + p_coeff_7 = lsq_rmat_rdiag_c_6 * (lsq_qtmat_c_6(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_6(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_6(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_6(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_6(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_6(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_6(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_6(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_6(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_4*p_coeff_8 + lsq_rmat_utri_c_5*p_coeff_9 + lsq_rmat_utri_c_6*p_coeff_10)) + + p_coeff_6 = lsq_rmat_rdiag_c_5 * (lsq_qtmat_c_5(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_5(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_5(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_5(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_5(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_5(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_5(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_5(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_5(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_7*p_coeff_7 + lsq_rmat_utri_c_8*p_coeff_8 + lsq_rmat_utri_c_9*p_coeff_9 + lsq_rmat_utri_c_10*p_coeff_10)) + + p_coeff_5 = lsq_rmat_rdiag_c_4 * (lsq_qtmat_c_4(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_4(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_4(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_4(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_4(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_4(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_4(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_4(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_4(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_11*p_coeff_6 + lsq_rmat_utri_c_12*p_coeff_7 + lsq_rmat_utri_c_13*p_coeff_8 + lsq_rmat_utri_c_14*p_coeff_9 + lsq_rmat_utri_c_15*p_coeff_10)) + + p_coeff_4 = lsq_rmat_rdiag_c_3 * (lsq_qtmat_c_3(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_3(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_3(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_3(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_3(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_3(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_3(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_3(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_3(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_16*p_coeff_5 + lsq_rmat_utri_c_17*p_coeff_6 + lsq_rmat_utri_c_18*p_coeff_7 + lsq_rmat_utri_c_19*p_coeff_8 + lsq_rmat_utri_c_20*p_coeff_9 + lsq_rmat_utri_c_21*p_coeff_10)) + + p_coeff_3 = lsq_rmat_rdiag_c_2 * (lsq_qtmat_c_2(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_2(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_2(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_2(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_2(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_2(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_2(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_2(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_2(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_22*p_coeff_4 + lsq_rmat_utri_c_23*p_coeff_5 + lsq_rmat_utri_c_24*p_coeff_6 + lsq_rmat_utri_c_25*p_coeff_7 + lsq_rmat_utri_c_26*p_coeff_8 + lsq_rmat_utri_c_27*p_coeff_9 + lsq_rmat_utri_c_28*p_coeff_10)) + + p_coeff_2 = lsq_rmat_rdiag_c_1 * (lsq_qtmat_c_1(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_1(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_1(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_1(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_1(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_1(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_1(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_1(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_1(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_29*p_coeff_3 + lsq_rmat_utri_c_30*p_coeff_4 + lsq_rmat_utri_c_31*p_coeff_5 + lsq_rmat_utri_c_32*p_coeff_6 + lsq_rmat_utri_c_33*p_coeff_7 + lsq_rmat_utri_c_34*p_coeff_8 + lsq_rmat_utri_c_35*p_coeff_9 + lsq_rmat_utri_c_36*p_coeff_10)) + + p_coeff_1 = p_cc - (p_coeff_2*lsq_moments_1 + p_coeff_3*lsq_moments_2 + p_coeff_4*lsq_moments_3 + p_coeff_5*lsq_moments_4 + p_coeff_6*lsq_moments_5 + p_coeff_7*lsq_moments_6 + p_coeff_8*lsq_moments_7 + p_coeff_9*lsq_moments_8 + p_coeff_10*lsq_moments_9) + return p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10 + +@program +def recon_lsq_cell_c_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_qtmat_c_1: Field[[CECECDim], float], + lsq_qtmat_c_2: Field[[CECECDim], float], + lsq_qtmat_c_3: Field[[CECECDim], float], + lsq_qtmat_c_4: Field[[CECECDim], float], + lsq_qtmat_c_5: Field[[CECECDim], float], + lsq_qtmat_c_6: Field[[CECECDim], float], + lsq_qtmat_c_7: Field[[CECECDim], float], + lsq_qtmat_c_8: Field[[CECECDim], float], + lsq_qtmat_c_9: Field[[CECECDim], float], + lsq_rmat_rdiag_c_1: Field[[CellDim], float], + lsq_rmat_rdiag_c_2: Field[[CellDim], float], + lsq_rmat_rdiag_c_3: Field[[CellDim], float], + lsq_rmat_rdiag_c_4: Field[[CellDim], float], + lsq_rmat_rdiag_c_5: Field[[CellDim], float], + lsq_rmat_rdiag_c_6: Field[[CellDim], float], + lsq_rmat_rdiag_c_7: Field[[CellDim], float], + lsq_rmat_rdiag_c_8: Field[[CellDim], float], + lsq_rmat_rdiag_c_9: Field[[CellDim], float], + lsq_rmat_utri_c_1: Field[[CellDim], float], + lsq_rmat_utri_c_2: Field[[CellDim], float], + lsq_rmat_utri_c_3: Field[[CellDim], float], + lsq_rmat_utri_c_4: Field[[CellDim], float], + lsq_rmat_utri_c_5: Field[[CellDim], float], + lsq_rmat_utri_c_6: Field[[CellDim], float], + lsq_rmat_utri_c_7: Field[[CellDim], float], + lsq_rmat_utri_c_8: Field[[CellDim], float], + lsq_rmat_utri_c_9: Field[[CellDim], float], + lsq_rmat_utri_c_10: Field[[CellDim], float], + lsq_rmat_utri_c_11: Field[[CellDim], float], + lsq_rmat_utri_c_12: Field[[CellDim], float], + lsq_rmat_utri_c_13: Field[[CellDim], float], + lsq_rmat_utri_c_14: Field[[CellDim], float], + lsq_rmat_utri_c_15: Field[[CellDim], float], + lsq_rmat_utri_c_16: Field[[CellDim], float], + lsq_rmat_utri_c_17: Field[[CellDim], float], + lsq_rmat_utri_c_18: Field[[CellDim], float], + lsq_rmat_utri_c_19: Field[[CellDim], float], + lsq_rmat_utri_c_20: Field[[CellDim], float], + lsq_rmat_utri_c_21: Field[[CellDim], float], + lsq_rmat_utri_c_22: Field[[CellDim], float], + lsq_rmat_utri_c_23: Field[[CellDim], float], + lsq_rmat_utri_c_24: Field[[CellDim], float], + lsq_rmat_utri_c_25: Field[[CellDim], float], + lsq_rmat_utri_c_26: Field[[CellDim], float], + lsq_rmat_utri_c_27: Field[[CellDim], float], + lsq_rmat_utri_c_28: Field[[CellDim], float], + lsq_rmat_utri_c_29: Field[[CellDim], float], + lsq_rmat_utri_c_30: Field[[CellDim], float], + lsq_rmat_utri_c_31: Field[[CellDim], float], + lsq_rmat_utri_c_32: Field[[CellDim], float], + lsq_rmat_utri_c_33: Field[[CellDim], float], + lsq_rmat_utri_c_34: Field[[CellDim], float], + lsq_rmat_utri_c_35: Field[[CellDim], float], + lsq_rmat_utri_c_36: Field[[CellDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], + p_coeff_1: Field[[CellDim, KDim], float], + p_coeff_2: Field[[CellDim, KDim], float], + p_coeff_3: Field[[CellDim, KDim], float], + p_coeff_4: Field[[CellDim, KDim], float], + p_coeff_5: Field[[CellDim, KDim], float], + p_coeff_6: Field[[CellDim, KDim], float], + p_coeff_7: Field[[CellDim, KDim], float], + p_coeff_8: Field[[CellDim, KDim], float], + p_coeff_9: Field[[CellDim, KDim], float], + p_coeff_10: Field[[CellDim, KDim], float], +): + _recon_lsq_cell_c_stencil(p_cc, lsq_qtmat_c_1, lsq_qtmat_c_2, lsq_qtmat_c_3, lsq_qtmat_c_4, lsq_qtmat_c_5, lsq_qtmat_c_6, lsq_qtmat_c_7, lsq_qtmat_c_8, lsq_qtmat_c_9, lsq_rmat_rdiag_c_1, lsq_rmat_rdiag_c_2, lsq_rmat_rdiag_c_3, lsq_rmat_rdiag_c_4, lsq_rmat_rdiag_c_5, lsq_rmat_rdiag_c_6, lsq_rmat_rdiag_c_7, lsq_rmat_rdiag_c_8, lsq_rmat_rdiag_c_9, lsq_rmat_utri_c_1, lsq_rmat_utri_c_2, lsq_rmat_utri_c_3, lsq_rmat_utri_c_4, lsq_rmat_utri_c_5, lsq_rmat_utri_c_6, lsq_rmat_utri_c_7, lsq_rmat_utri_c_8, lsq_rmat_utri_c_9, lsq_rmat_utri_c_10, lsq_rmat_utri_c_11, lsq_rmat_utri_c_12, lsq_rmat_utri_c_13, lsq_rmat_utri_c_14, lsq_rmat_utri_c_15, lsq_rmat_utri_c_16, lsq_rmat_utri_c_17, lsq_rmat_utri_c_18, lsq_rmat_utri_c_19, lsq_rmat_utri_c_20, lsq_rmat_utri_c_21, lsq_rmat_utri_c_22, lsq_rmat_utri_c_23, lsq_rmat_utri_c_24, lsq_rmat_utri_c_25, lsq_rmat_utri_c_26, lsq_rmat_utri_c_27, lsq_rmat_utri_c_28, lsq_rmat_utri_c_29, lsq_rmat_utri_c_30, lsq_rmat_utri_c_31, lsq_rmat_utri_c_32, lsq_rmat_utri_c_33, lsq_rmat_utri_c_34, lsq_rmat_utri_c_35, lsq_rmat_utri_c_36, lsq_moments_1, lsq_moments_2, lsq_moments_3, lsq_moments_4, lsq_moments_5, lsq_moments_6, lsq_moments_7, lsq_moments_8, lsq_moments_9, out=(p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10) + ) From f861a3f39147b44698d7ffdd1b83dd9a8b355f96 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 13 Jan 2023 16:00:06 +0100 Subject: [PATCH 013/105] CECEC + C2E2C2E2C definitions added --- common/src/icon4py/common/dimension.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 53f7904193..6b1c5d2adf 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -23,6 +23,7 @@ ECDim = Dimension("EC") ECVDim = Dimension("ECV") CECDim = Dimension("CEC") +CECECDim = Dimension("CECEC") E2CDim = Dimension("E2C", DimensionKind.LOCAL) E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) @@ -33,6 +34,7 @@ E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL) E2C2EDim = Dimension("E2C2E", DimensionKind.LOCAL) C2E2CDim = Dimension("C2E2C", DimensionKind.LOCAL) +C2E2C2E2CDim = Dimension("C2E2C2E2C", DimensionKind.LOCAL) E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim)) C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim)) V2C = FieldOffset("V2C", source=CellDim, target=(VertexDim, V2CDim)) @@ -46,6 +48,8 @@ E2C2EO = FieldOffset("E2C2EO", source=EdgeDim, target=(EdgeDim, E2C2EODim)) E2C2E = FieldOffset("E2C2E", source=EdgeDim, target=(EdgeDim, E2C2EDim)) C2E2C = FieldOffset("C2E2C", source=CellDim, target=(CellDim, C2E2CDim)) +C2E2C2E2C = FieldOffset("C2E2C2E2C", source=CellDim, target=(CellDim, C2E2C2E2CDim)) C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim)) +C2CECEC = FieldOffset("C2CECEC", source=CECECDim, target=(CellDim, C2E2C2E2CDim)) Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) From 9af4c877d8aa8c3825a9c25ef9802e98c5d15e83 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 9 Feb 2023 14:33:39 +0100 Subject: [PATCH 014/105] first try --- .../advection/hflx_limiter_mo_stencil_01a.py | 51 +++++++++++++ .../advection/hflx_limiter_mo_stencil_01b.py | 71 +++++++++++++++++++ 2 files changed, 122 insertions(+) create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py new file mode 100644 index 0000000000..727dcd83bc --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py @@ -0,0 +1,51 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + + z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) - abs(p_mass_flx_e)* ( p_cc(E2C[0]) - p_cc(E2C[1]) ) ) + # z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) ) + z_anti = p_mflx_tracer_h - z_mflx_low + + return (z_anti, z_mflx_low) + + +@program +def hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h, + p_mass_flx_e, + z_mflx_low, + z_anti, + p_cc, + out=(z_anti, z_mflx_low), + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py new file mode 100644 index 0000000000..689e6dd8ad --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py @@ -0,0 +1,71 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, neighbor_sum + +from icon4py.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim, C2EDim, E2C + + +@field_operator +def _hflx_limiter_mo_stencil_01b( + geofac_div: Field[[CellDim, C2EDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_rhodz_new: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_fluxdiv_c: Field[[CellDim, KDim], float], + p_dtime: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + z_mflx_anti = broadcast(p_rhodz_new, (CellDim, C2EDim, KDim)) + + z_mflx_anti_in = -1. * neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) + z_mflx_anti_out = neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) + z_fluxdiv_c = neighbor_sum( z_mflx_low(C2E) * geofac_div, axis=C2EDim) + + z_tracer_new_low = ( p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c ) / p_rhodz_new + z_tracer_max = maximum(p_cc,z_tracer_new_low) + z_tracer_min = minimum(p_cc,z_tracer_new_low) + + return (z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, z_tracer_max, z_tracer_min) + + +@program +def hflx_limiter_mo_stencil_01b( + geofac_div: Field[[CellDim, C2EDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_rhodz_new: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_fluxdiv_c: Field[[CellDim, KDim], float], + p_dtime: float, + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_01b( + geofac_div, + p_rhodz_now, + p_rhodz_new, + z_mflx_low, + z_anti, + p_cc, + z_fluxdiv_c, + p_dtime, + out=(z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, z_tracer_max, z_tracer_min), + ) From 6f5b7b492bd77ae146d06ce54289e665fb6a5a5f Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 Feb 2023 12:25:30 +0100 Subject: [PATCH 015/105] add _face_val_ppm_stencil_01 --- .../advection/face_val_ppm_stencil_01a.py | 43 +++++++++++++++++++ .../advection/face_val_ppm_stencil_01b.py | 42 ++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_01a.py create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_01b.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01a.py new file mode 100644 index 0000000000..fcb2624992 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_01a.py @@ -0,0 +1,43 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_01a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) + zfac = (p_cc(Koff[+1]) - p_cc) / (p_cellhgt_mc_now(Koff[+1]) + p_cellhgt_mc_now) + z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[+1])) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now(Koff[+1])) * zfac_m1) + + return z_slope + + +@program +def face_val_ppm_stencil_01a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_01a( + p_cc, + p_cellhgt_mc_now, + out=z_slope, + ) \ No newline at end of file diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py new file mode 100644 index 0000000000..05f89c5529 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py @@ -0,0 +1,42 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_01b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) + zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) + z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now) * zfac_m1) + return z_slope + + +@program +def face_val_ppm_stencil_01b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_01b( + p_cc, + p_cellhgt_mc_now, + out=z_slope, + ) \ No newline at end of file From ccce8594098fd0270e2216ca5be899bf0cfd9bb7 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 Feb 2023 13:31:39 +0100 Subject: [PATCH 016/105] add more stencils --- .../advection/face_val_ppm_stencil_03a.py | 40 +++++++++++++++++++ .../advection/face_val_ppm_stencil_03b.py | 40 +++++++++++++++++++ .../advection/face_val_ppm_stencil_03c.py | 37 +++++++++++++++++ 3 files changed, 117 insertions(+) create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03a.py create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03b.py create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03c.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py new file mode 100644 index 0000000000..d2aaadab37 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py @@ -0,0 +1,40 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_03a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc*(1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))* p_cc + p_cc(Koff[-1])) + return p_face + + +@program +def face_val_ppm_stencil_03a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_03a( + p_cc, + p_cellhgt_mc_now, + out=p_face, + ) \ No newline at end of file diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py new file mode 100644 index 0000000000..21608be3b5 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py @@ -0,0 +1,40 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_03b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc*( 1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) + return p_face + + +@program +def face_val_ppm_stencil_03b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_03b( + p_cc, + p_cellhgt_mc_now, + out=p_face, + ) \ No newline at end of file diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py new file mode 100644 index 0000000000..938bb63d68 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py @@ -0,0 +1,37 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_03c( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc + return p_face + + +@program +def face_val_ppm_stencil_03c( + p_cc: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_03c( + p_cc, + out=p_face, + ) \ No newline at end of file From 090147b376f8d28b70c9225252a6f505be8f8083 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 16 Feb 2023 16:11:21 +0100 Subject: [PATCH 017/105] rename file --- .../{face_val_ppm_stencil_01a.py => face_val_ppm_stencil_01.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename advection/src/icon4py/advection/{face_val_ppm_stencil_01a.py => face_val_ppm_stencil_01.py} (100%) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_01a.py rename to advection/src/icon4py/advection/face_val_ppm_stencil_01.py From 1c781e412072dad183c1d2c5c17676b05eaf4c31 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 16 Feb 2023 17:11:49 +0100 Subject: [PATCH 018/105] merge stencils --- .../advection/face_val_ppm_stencil_01.py | 40 +++++++++++++++++-- .../advection/face_val_ppm_stencil_01b.py | 1 + 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py index fcb2624992..c78e704b1e 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_01.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, broadcast, maximum, minimum, abs +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs, int32, where from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff @@ -30,14 +30,48 @@ def _face_val_ppm_stencil_01a( return z_slope +@field_operator +def _face_val_ppm_stencil_01b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) + zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) + z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now) * zfac_m1) + + return z_slope + + +@field_operator +def _face_val_ppm_stencil_01( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + elev: int32, +) -> Field[[CellDim, KDim], float]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + z_slope = where(vert_idx Date: Thu, 16 Feb 2023 18:12:40 +0100 Subject: [PATCH 019/105] add stencil 2 --- .../advection/face_val_ppm_stencil_02.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_02.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py new file mode 100644 index 0000000000..c349247460 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py @@ -0,0 +1,93 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs, int32, where + +from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc*(1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))* p_cc + p_cc(Koff[-1])) + + return p_face + + +@field_operator +def _face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc + return p_face + +@field_operator +def _face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc(Koff[-1]) + return p_face + + +@field_operator +def _face_val_ppm_stencil_02( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face_in: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32 +) -> Field[[CellDim, KDim], float]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + p_face = where( (vert_idx==slevp1) | (vert_idx==elev), _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), p_face_in) + + p_face = where( (vert_idx==slev), _face_val_ppm_stencil_02b(p_cc), p_face_in) + + p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02b(p_cc), p_face_in) + + return p_face + + +@program +def face_val_ppm_stencil_02( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face_in: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02( + p_cc, + p_cellhgt_mc_now, + p_face_in, + vert_idx, + slev, + elev, + slevp1, + elevp1, + out=p_face, + ) \ No newline at end of file From b2e853663df72bb3554dad35be9de0ba13e4704e Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 16 Feb 2023 18:14:32 +0100 Subject: [PATCH 020/105] remove stencil 3 --- .../advection/face_val_ppm_stencil_03a.py | 40 ------------------- .../advection/face_val_ppm_stencil_03b.py | 40 ------------------- .../advection/face_val_ppm_stencil_03c.py | 37 ----------------- 3 files changed, 117 deletions(-) delete mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03a.py delete mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03b.py delete mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_03c.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py deleted file mode 100644 index d2aaadab37..0000000000 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_03a.py +++ /dev/null @@ -1,40 +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.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs - -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff - - -@field_operator -def _face_val_ppm_stencil_03a( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], -) -> Field[[CellDim, KDim], float]: - - p_face = p_cc*(1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))* p_cc + p_cc(Koff[-1])) - return p_face - - -@program -def face_val_ppm_stencil_03a( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], - p_face: Field[[CellDim, KDim], float], -): - _face_val_ppm_stencil_03a( - p_cc, - p_cellhgt_mc_now, - out=p_face, - ) \ No newline at end of file diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py deleted file mode 100644 index 21608be3b5..0000000000 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_03b.py +++ /dev/null @@ -1,40 +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.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs - -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff - - -@field_operator -def _face_val_ppm_stencil_03b( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], -) -> Field[[CellDim, KDim], float]: - - p_face = p_cc*( 1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) - return p_face - - -@program -def face_val_ppm_stencil_03b( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], - p_face: Field[[CellDim, KDim], float], -): - _face_val_ppm_stencil_03b( - p_cc, - p_cellhgt_mc_now, - out=p_face, - ) \ No newline at end of file diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py b/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py deleted file mode 100644 index 938bb63d68..0000000000 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_03c.py +++ /dev/null @@ -1,37 +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.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs - -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff - - -@field_operator -def _face_val_ppm_stencil_03c( - p_cc: Field[[CellDim, KDim], float], -) -> Field[[CellDim, KDim], float]: - - p_face = p_cc - return p_face - - -@program -def face_val_ppm_stencil_03c( - p_cc: Field[[CellDim, KDim], float], - p_face: Field[[CellDim, KDim], float], -): - _face_val_ppm_stencil_03c( - p_cc, - out=p_face, - ) \ No newline at end of file From 2f018e2d8f19f2826ba4dcb253771538578e3fd1 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 22 Feb 2023 13:22:48 +0100 Subject: [PATCH 021/105] fix bounds --- .../tests/test_vlimit_prbl_sm_stencil_01.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 7367d96212..9e53f46c1f 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -25,23 +25,22 @@ def v_limit_prbl_sm_stencil_01_numpy( p_face: np.array, p_cc: np.array, ): - p_face_k_plus_1 = np.roll(p_face, shift=-1, axis=1) - z_delta = p_face - p_face_k_plus_1 - z_a6i = 6.0 * (p_cc - 0.5 * (p_face + p_face_k_plus_1)) + z_delta = p_face[:,:-1] - p_face[:,1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:,:-1] + p_face[:,1:])) q_face_up, q_face_low = np.where( abs(z_delta) < -1 * z_a6i, np.where( - (p_cc < np.minimum(p_face, p_face_k_plus_1)), + (p_cc < np.minimum(p_face[:,:-1], p_face[:,1:])), (p_cc, p_cc), np.where( - p_face > p_face_k_plus_1, - (3.0 * p_cc - 2.0 * p_face_k_plus_1, p_face_k_plus_1), - (p_face, 3.0 * p_cc - 2.0 * p_face), + p_face[:,:-1] > p_face[:,1:], + (3.0 * p_cc - 2.0 * p_face[:,1:], p_face[:,1:]), + (p_face[:,:-1], 3.0 * p_cc - 2.0 * p_face[:,:-1]), ), ), - (p_face, p_face_k_plus_1), + (p_face[:,:-1], p_face[:,1:]), ) return q_face_up, q_face_low @@ -50,7 +49,7 @@ def v_limit_prbl_sm_stencil_01_numpy( def test_v_limit_prbl_sm_stencil_01(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) - p_face = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) p_face_up = zero_field(mesh, CellDim, KDim) p_face_low = zero_field(mesh, CellDim, KDim) From 11e63e315afc70c62ff5ac6ea47ef2d3e34d67a5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 22 Feb 2023 13:36:12 +0100 Subject: [PATCH 022/105] fix test --- advection/tests/test_vert_adv_stencil_01.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py index c8de7c3dc2..4f034feed8 100644 --- a/advection/tests/test_vert_adv_stencil_01.py +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -28,12 +28,11 @@ def vert_adv_stencil_01_numpy( rhodz_new: np.array, p_dtime, ) -> np.array: - p_mflx_tracer_v_offset_1 = np.roll(p_mflx_tracer_v, shift=-1, axis=1) tracer_new = ( tracer_now * rhodz_now + p_dtime - * (p_mflx_tracer_v_offset_1 * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + * (p_mflx_tracer_v[:,1:] * deepatmo_divzl - p_mflx_tracer_v[:,:-1] * deepatmo_divzu) ) / rhodz_new return tracer_new @@ -44,7 +43,7 @@ def test_vert_adv_stencil_01(): tracer_now = random_field(mesh, CellDim, KDim) rhodz_now = random_field(mesh, CellDim, KDim) - p_mflx_tracer_v = random_field(mesh, CellDim, KDim) + p_mflx_tracer_v = random_field(mesh, CellDim, KDim, extend={KDim: 1}) deepatmo_divzl = random_field(mesh, KDim) deepatmo_divzu = random_field(mesh, KDim) rhodz_new = random_field(mesh, CellDim, KDim) From 4ec708bd76d72d6b593683e66cfeb89e844ab053 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 22 Feb 2023 13:43:47 +0100 Subject: [PATCH 023/105] fix test --- advection/tests/test_step_advection_stencil_02.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py index a9a5e99ca9..2aff78df0e 100644 --- a/advection/tests/test_step_advection_stencil_02.py +++ b/advection/tests/test_step_advection_stencil_02.py @@ -28,8 +28,8 @@ def step_advection_stencil_02_numpy( ) -> np.ndarray: tmp = ( - np.roll(p_mflx_contra_v, axis=1, shift=-1) * deepatmo_divzl - - p_mflx_contra_v * deepatmo_divzu + p_mflx_contra_v[:,1:] * deepatmo_divzl + - p_mflx_contra_v[:,:-1] * deepatmo_divzu ) return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp @@ -37,7 +37,7 @@ def step_advection_stencil_02_numpy( def test_step_advection_stencil_02(): mesh = SimpleMesh() rhodz_ast = random_field(mesh, CellDim, KDim) - p_mflx_contra = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim, extend={KDim: 1}) deepatmo_divzl = random_field(mesh, KDim) deepatmo_divzu = random_field(mesh, KDim) result = zero_field(mesh, CellDim, KDim) From 477320125da1b3b1d69255c5cc8885430b7d7d28 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 22 Feb 2023 14:32:34 +0100 Subject: [PATCH 024/105] fix test --- advection/tests/test_face_val_ppm_stencil_05.py | 17 ++++++++++------- .../tests/test_step_advection_stencil_01.py | 5 ++--- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/advection/tests/test_face_val_ppm_stencil_05.py index c36690a5cf..8a85d8e088 100644 --- a/advection/tests/test_face_val_ppm_stencil_05.py +++ b/advection/tests/test_face_val_ppm_stencil_05.py @@ -24,12 +24,15 @@ def face_val_ppm_stencil_05_numpy( p_cellhgt_mc_now: np.array, z_slope: np.array, ): - p_cellhgt_mc_now_k_minus_1 = np.roll(p_cellhgt_mc_now, shift=1, axis=1) - p_cellhgt_mc_now_k_minus_2 = np.roll(p_cellhgt_mc_now, shift=2, axis=1) - p_cellhgt_mc_now_k_plus_1 = np.roll(p_cellhgt_mc_now, shift=-1, axis=1) + p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:,1:-2] + p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:,0:-3] + p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:,3:] + p_cellhgt_mc_now = p_cellhgt_mc_now[:,2:-1] - p_cc_k_minus_1 = np.roll(p_cc, shift=1, axis=1) - z_slope_k_minus_1 = np.roll(z_slope, shift=1, axis=1) + p_cc_k_minus_1 = p_cc[:,1:-1] + p_cc = p_cc[:,2:] + z_slope_k_minus_1 = z_slope[:,1:-1] + z_slope = z_slope[:,2:] zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) zgeo2 = 1.0 / ( @@ -62,7 +65,7 @@ def face_val_ppm_stencil_05_numpy( def test_face_val_ppm_stencil_05(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) - p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim, extend={KDim: 1}) z_slope = random_field(mesh, CellDim, KDim) p_face = zero_field(mesh, CellDim, KDim) @@ -80,4 +83,4 @@ def test_face_val_ppm_stencil_05(): offset_provider={"Koff": KDim}, ) - assert np.allclose(ref[:, 1:-1], p_face.__array__()[:, 1:-1]) + assert np.allclose(ref[:,:], p_face[:, 2:]) diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py index fc2290de19..9c2621e4a4 100644 --- a/advection/tests/test_step_advection_stencil_01.py +++ b/advection/tests/test_step_advection_stencil_01.py @@ -26,9 +26,8 @@ def step_advection_stencil_01_numpy( deepatmo_divzu: np.ndarray, pd_time: float, ) -> np.ndarray: - p_mlfx_contra_v1 = np.roll(p_mflx_contra_v, axis=1, shift=-1) tmp = pd_time * ( - p_mlfx_contra_v1 * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + p_mflx_contra_v[:,1:] * deepatmo_divzl - p_mflx_contra_v[:,:-1] * deepatmo_divzu ) return rhodz_ast + tmp @@ -36,7 +35,7 @@ def step_advection_stencil_01_numpy( def test_step_advection_stencil_01(): mesh = SimpleMesh() rhodz_ast = random_field(mesh, CellDim, KDim) - p_mflx_contra = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim, extend={KDim: 1}) deepatmo_divzl = random_field(mesh, KDim) deepatmo_divzu = random_field(mesh, KDim) result = zero_field(mesh, CellDim, KDim) From bf80cf486664caeeebf0d3a9c635c53b9ebff6e3 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 23 Feb 2023 12:27:39 +0100 Subject: [PATCH 025/105] add test --- .../tests/test_face_val_ppm_stencil_01.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 advection/tests/test_face_val_ppm_stencil_01.py diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py new file mode 100644 index 0000000000..345d36f0c5 --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -0,0 +1,85 @@ +# 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 + +import numpy as np + +from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field +from gt4py.next.ffront.fbuiltins import int32 + + +def face_val_ppm_stencil_01_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + vert_idx: np.array, + elev: int32, + z_slope: np.array, +): + + # vert_idx = np.broadcast_to(vert_idx, p_cc.shape) + + # 01a + zfac_m1 = (p_cc[:,1:-1] - p_cc[:,:-2]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,:-2]) + zfac = (p_cc[:,2:] - p_cc[:,1:-1]) / (p_cellhgt_mc_now[:,2:] + p_cellhgt_mc_now[:,1:-1]) + z_slope_a = ( p_cellhgt_mc_now[:,1:-1] / (p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,2:]) ) * ( (2.*p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1]) * zfac + (p_cellhgt_mc_now[:,1:-1] + 2.*p_cellhgt_mc_now[:,2:]) * zfac_m1) + + + # 01b + zfac_m1 = (p_cc[:,1:-1] - p_cc[:,:-2]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,:-2]) + zfac = (p_cc[:,1:-1] - p_cc[:,1:-1]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,1:-1]) + z_slope_b = ( p_cellhgt_mc_now[:,1:-1] / (p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,1:-1]) ) * ( (2.*p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1]) * zfac + (p_cellhgt_mc_now[:,1:-1] + 2.*p_cellhgt_mc_now[:,1:-1]) * zfac_m1) + + z_slope = np.where(vert_idx[1:-1] Date: Thu, 23 Feb 2023 13:23:00 +0100 Subject: [PATCH 026/105] cleanup --- advection/tests/test_face_val_ppm_stencil_01.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 345d36f0c5..0c7c7a572a 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -59,9 +59,6 @@ def test_face_val_ppm_stencil_01(): vert_idx[i] = i elev = vert_idx[-2] - # print(vert_idx.__array__()) - # print(elev) - z_slope = random_field(mesh, CellDim, KDim) ref = face_val_ppm_stencil_01_numpy( @@ -81,5 +78,4 @@ def test_face_val_ppm_stencil_01(): offset_provider={"Koff": KDim}, ) - # assert np.allclose(ref[:,1:], z_slope[:,1:-1]) assert np.allclose(ref[:,:-1], z_slope[:,1:-1]) From 9e32f4ff1adf24bdc6d760d6e8058ba7342c89c6 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 23 Feb 2023 17:44:13 +0100 Subject: [PATCH 027/105] fix --- advection/src/icon4py/advection/face_val_ppm_stencil_02.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py index c349247460..2477d928fd 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py @@ -63,7 +63,7 @@ def _face_val_ppm_stencil_02( p_face = where( (vert_idx==slev), _face_val_ppm_stencil_02b(p_cc), p_face_in) - p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02b(p_cc), p_face_in) + p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02c(p_cc), p_face_in) return p_face From d88e8995bf25455a1c3f1ceca5298c7cdba2b02b Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 23 Feb 2023 18:32:12 +0100 Subject: [PATCH 028/105] fix stencil --- advection/src/icon4py/advection/face_val_ppm_stencil_02.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py index 2477d928fd..2f40ccf840 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02.py @@ -61,9 +61,9 @@ def _face_val_ppm_stencil_02( p_face = where( (vert_idx==slevp1) | (vert_idx==elev), _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), p_face_in) - p_face = where( (vert_idx==slev), _face_val_ppm_stencil_02b(p_cc), p_face_in) + p_face = where( (vert_idx==slev), _face_val_ppm_stencil_02b(p_cc), p_face) - p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02c(p_cc), p_face_in) + p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) return p_face From 25f16ee136be475eadd0391f38f00e8a0e9c56dc Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 24 Feb 2023 11:18:43 +0100 Subject: [PATCH 029/105] cleanup --- advection/tests/test_face_val_ppm_stencil_01.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 0c7c7a572a..6feba84896 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -25,7 +25,6 @@ def face_val_ppm_stencil_01_numpy( p_cellhgt_mc_now: np.array, vert_idx: np.array, elev: int32, - z_slope: np.array, ): # vert_idx = np.broadcast_to(vert_idx, p_cc.shape) @@ -66,7 +65,6 @@ def test_face_val_ppm_stencil_01(): np.asarray(p_cellhgt_mc_now), np.asarray(vert_idx), elev, - np.asarray(z_slope), ) face_val_ppm_stencil_01( From 9a7c47545cc2bab673531330e9a50874eda361fb Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 24 Feb 2023 12:20:04 +0100 Subject: [PATCH 030/105] add second test --- .../tests/test_face_val_ppm_stencil_02.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 advection/tests/test_face_val_ppm_stencil_02.py diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py new file mode 100644 index 0000000000..18a21c5abd --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -0,0 +1,93 @@ +# 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 + +import numpy as np + +from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field +from gt4py.next.ffront.fbuiltins import int32 + + +def face_val_ppm_stencil_02_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + vert_idx: np.array, + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, +): + + p_face_a = p_face_in + + p_face_a[:,1:] = p_cc[:,1:]*(1. - (p_cellhgt_mc_now[:,1:] / p_cellhgt_mc_now[:,:-1])) + (p_cellhgt_mc_now[:,1:]/(p_cellhgt_mc_now[:,:-1] + p_cellhgt_mc_now[:,1:])) * ((p_cellhgt_mc_now[:,1:] / p_cellhgt_mc_now[:,:-1])* p_cc[:,1:] + p_cc[:,:-1]) + + p_face = np.where( (vert_idx==slevp1) | (vert_idx==elev), p_face_a, p_face_in) + p_face = np.where( (vert_idx==slev), p_cc, p_face) + p_face[:,1:] = np.where( (vert_idx[1:]==elevp1), p_cc[:,:-1], p_face[:,1:]) + + return p_face + + + +def test_face_val_ppm_stencil_02(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + p_face_in = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + vert_idx = zero_field(mesh, KDim, dtype=int32) + + for i in range(len(vert_idx.__array__())): + vert_idx[i] = i + + + slev = int32(1) + slevp1 = slev + int32(1) + elev = vert_idx[-3] + elevp1 = elev + int32(1) + + # print(vert_idx.__array__()) + # print(elev) + + ref = face_val_ppm_stencil_02_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + np.asarray(p_face_in), + np.asarray(vert_idx), + slev, + elev, + slevp1, + elevp1, + ) + + face_val_ppm_stencil_02( + p_cc, + p_cellhgt_mc_now, + p_face_in, + vert_idx, + slev, + elev, + slevp1, + elevp1, + p_face, + offset_provider={"Koff": KDim}, + ) + + # assert np.allclose(ref[:,1:], p_face) + # assert np.allclose(ref, p_face[:,1:]) + # assert np.allclose(ref, p_face[:,1:]) + assert np.allclose(ref, p_face[:,:]) From f31ce798e9a9d3391f072281145219b58eaaf76c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 24 Feb 2023 13:59:35 +0100 Subject: [PATCH 031/105] cleanup --- advection/tests/test_face_val_ppm_stencil_01.py | 7 ++++--- advection/tests/test_face_val_ppm_stencil_02.py | 13 ++++--------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 6feba84896..32548af981 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -16,8 +16,10 @@ from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.common.dimension import CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import random_field, zero_field, _shape from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + def face_val_ppm_stencil_01_numpy( @@ -54,8 +56,7 @@ def test_face_val_ppm_stencil_01(): p_cellhgt_mc_now = random_field(mesh, CellDim, KDim, extend={KDim: 1}) vert_idx = zero_field(mesh, KDim, dtype=int32, extend={KDim: 1}) - for i in range(len(vert_idx.__array__())): - vert_idx[i] = i + vert_idx = it_embedded.np_as_located_field(KDim)( np.arange(0, _shape(mesh, KDim, extend={KDim: 1})[0], dtype=int32) ) elev = vert_idx[-2] z_slope = random_field(mesh, CellDim, KDim) diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py index 18a21c5abd..1d1b0077a2 100644 --- a/advection/tests/test_face_val_ppm_stencil_02.py +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -16,8 +16,9 @@ from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.common.dimension import CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import random_field, zero_field, _shape from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded def face_val_ppm_stencil_02_numpy( @@ -49,12 +50,9 @@ def test_face_val_ppm_stencil_02(): p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) p_face_in = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim) - vert_idx = zero_field(mesh, KDim, dtype=int32) - for i in range(len(vert_idx.__array__())): - vert_idx[i] = i + vert_idx = it_embedded.np_as_located_field(KDim)( np.arange(0, _shape(mesh, KDim)[0], dtype=int32) ) - slev = int32(1) slevp1 = slev + int32(1) elev = vert_idx[-3] @@ -87,7 +85,4 @@ def test_face_val_ppm_stencil_02(): offset_provider={"Koff": KDim}, ) - # assert np.allclose(ref[:,1:], p_face) - # assert np.allclose(ref, p_face[:,1:]) - # assert np.allclose(ref, p_face[:,1:]) - assert np.allclose(ref, p_face[:,:]) + assert np.allclose(ref, p_face) From ab9255d35be391fec0231e5ce9f0c1420893f9c4 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 24 Feb 2023 14:30:38 +0100 Subject: [PATCH 032/105] rm print --- advection/tests/test_face_val_ppm_stencil_02.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py index 1d1b0077a2..6473a499a1 100644 --- a/advection/tests/test_face_val_ppm_stencil_02.py +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -58,9 +58,6 @@ def test_face_val_ppm_stencil_02(): elev = vert_idx[-3] elevp1 = elev + int32(1) - # print(vert_idx.__array__()) - # print(elev) - ref = face_val_ppm_stencil_02_numpy( np.asarray(p_cc), np.asarray(p_cellhgt_mc_now), From 38520aa814314b337ae6b7131eb8d2265f53a05a Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 3 Mar 2023 15:12:57 +0100 Subject: [PATCH 033/105] fix stencil --- .../advection/hflx_limiter_mo_stencil_02.py | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py index 9222ff0216..2160275ee7 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py @@ -23,20 +23,22 @@ def _hflx_limiter_mo_stencil_02( refin_ctrl: Field[[CellDim], int32], p_cc: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], lo_bound: int32, hi_bound: int32, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound) - z_tracer_new_tmp = where( + z_tracer_new_out = where( condition, minimum(1.1 * p_cc, maximum(0.9 * p_cc, z_tracer_new_low)), z_tracer_new_low, ) - return where( - condition, - (maximum(p_cc, z_tracer_new_tmp), minimum(p_cc, z_tracer_new_tmp)), - (z_tracer_new_low, z_tracer_new_low), - ) + + z_tracer_max_out = where(condition, maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = where(condition, minimum(p_cc, z_tracer_new_out), z_tracer_min) + + return(z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) @program @@ -44,16 +46,21 @@ def hflx_limiter_mo_stencil_02( refin_ctrl: Field[[CellDim], int32], p_cc: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], - lo_bound: int32, - hi_bound: int32, z_tracer_max: Field[[CellDim, KDim], float], z_tracer_min: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, + z_tracer_new_low_out: Field[[CellDim, KDim], float], + z_tracer_max_out: Field[[CellDim, KDim], float], + z_tracer_min_out: Field[[CellDim, KDim], float], ): _hflx_limiter_mo_stencil_02( refin_ctrl, p_cc, z_tracer_new_low, + z_tracer_max, + z_tracer_min, lo_bound, hi_bound, - out=(z_tracer_max, z_tracer_min), + out=(z_tracer_new_low_out, z_tracer_max_out, z_tracer_min_out), ) From c6c66d4c3305af5c07635e57758cc1db17214e34 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 3 Mar 2023 15:55:19 +0100 Subject: [PATCH 034/105] fix test --- .../tests/test_hflx_limiter_mo_stencil_02.py | 60 ++++++++++++++----- 1 file changed, 46 insertions(+), 14 deletions(-) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py index 13ed590a59..e786bd4545 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -26,6 +26,8 @@ def hflx_limiter_mo_stencil_02_numpy( refin_ctrl: np.ndarray, p_cc: np.ndarray, z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, lo_bound: float, hi_bound: float, ): @@ -34,34 +36,45 @@ def hflx_limiter_mo_stencil_02_numpy( np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), ) - z_tracer_new_tmp = np.where( + z_tracer_new_out = np.where( condition, np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), z_tracer_new_low, ) - z_tracer_max = np.where( - condition, np.maximum(p_cc, z_tracer_new_tmp), z_tracer_new_low + z_tracer_max_out = np.where( + condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max ) - z_tracer_min = np.where( - condition, np.minimum(p_cc, z_tracer_new_tmp), z_tracer_new_low + z_tracer_min_out = np.where( + condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min ) - return z_tracer_max, z_tracer_min + return z_tracer_new_out, z_tracer_max_out, z_tracer_min_out def test_hflx_limiter_mo_stencil_02_some_matching_condition(): mesh = SimpleMesh() + hi_bound = np.int32(1) lo_bound = np.int32(5) + refin_ctrl = constant_field(mesh, hi_bound, CellDim, dtype=int32) + refin_ctrl[0:2] = np.int32(3) + p_cc = random_field(mesh, CellDim, KDim) - z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_new_low_in = random_field(mesh, CellDim, KDim) + z_tracer_max_in = random_field(mesh, CellDim, KDim) + z_tracer_min_in = random_field(mesh, CellDim, KDim) + + z_tracer_new_low = zero_field(mesh, CellDim, KDim) z_tracer_max = zero_field(mesh, CellDim, KDim) z_tracer_min = zero_field(mesh, CellDim, KDim) - ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( + + ref_new_low, ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( np.asarray(refin_ctrl), np.asarray(p_cc), - np.asarray(z_tracer_new_low), + np.asarray(z_tracer_new_low_in), + np.asarray(z_tracer_max_in), + np.asarray(z_tracer_min_in), lo_bound, hi_bound, ) @@ -69,35 +82,54 @@ def test_hflx_limiter_mo_stencil_02_some_matching_condition(): hflx_limiter_mo_stencil_02( refin_ctrl, p_cc, - z_tracer_new_low, + z_tracer_new_low_in, + z_tracer_max_in, + z_tracer_min_in, lo_bound, hi_bound, + z_tracer_new_low, z_tracer_max, z_tracer_min, offset_provider={}, ) + + assert np.allclose(z_tracer_new_low, ref_new_low) assert np.allclose(z_tracer_max, ref_max) assert np.allclose(z_tracer_min, ref_min) def test_hflx_limiter_mo_stencil_02_none_matching_condition(): mesh = SimpleMesh() + hi_bound = np.int32(3) lo_bound = np.int32(1) + refin_ctrl = constant_field(mesh, 2, CellDim, dtype=int32) + p_cc = random_field(mesh, CellDim, KDim) - z_tracer_new_low = random_field(mesh, CellDim, KDim) + + z_tracer_new_low_in = random_field(mesh, CellDim, KDim) + z_tracer_max_in = random_field(mesh, CellDim, KDim) + z_tracer_min_in = random_field(mesh, CellDim, KDim) + + z_tracer_new_low = zero_field(mesh, CellDim, KDim) z_tracer_max = zero_field(mesh, CellDim, KDim) z_tracer_min = zero_field(mesh, CellDim, KDim) + hflx_limiter_mo_stencil_02( refin_ctrl, p_cc, - z_tracer_new_low, + z_tracer_new_low_in, + z_tracer_max_in, + z_tracer_min_in, lo_bound, hi_bound, + z_tracer_new_low, z_tracer_max, z_tracer_min, offset_provider={}, ) - assert np.allclose(z_tracer_min, z_tracer_new_low) - assert np.allclose(z_tracer_max, z_tracer_new_low) + + assert np.allclose(z_tracer_new_low_in, z_tracer_new_low) + assert np.allclose(z_tracer_min_in, z_tracer_min) + assert np.allclose(z_tracer_max_in, z_tracer_max) From fc37993ea134c420d83889c7b806ee2037f21844 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 3 Mar 2023 16:01:41 +0100 Subject: [PATCH 035/105] rm stenc --- .../advection/face_val_ppm_stencil_01b.py | 43 ------------------- 1 file changed, 43 deletions(-) delete mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_01b.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py deleted file mode 100644 index 63eff80194..0000000000 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_01b.py +++ /dev/null @@ -1,43 +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.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs - -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff - - -@field_operator -def _face_val_ppm_stencil_01b( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], -) -> Field[[CellDim, KDim], float]: - - zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) - zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) - z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now) * zfac_m1) - - return z_slope - - -@program -def face_val_ppm_stencil_01b( - p_cc: Field[[CellDim, KDim], float], - p_cellhgt_mc_now: Field[[CellDim, KDim], float], - z_slope: Field[[CellDim, KDim], float], -): - _face_val_ppm_stencil_01b( - p_cc, - p_cellhgt_mc_now, - out=z_slope, - ) \ No newline at end of file From 976f0de903a9dbd74e30ead8aad3d0df972d2d55 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 10 Mar 2023 13:53:23 +0100 Subject: [PATCH 036/105] WIP stencil 03 --- .../advection/hflx_limiter_mo_stencil_03.py | 50 +++++++++++++++---- .../tests/test_hflx_limiter_mo_stencil_03.py | 8 +-- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py index 82ec3da2a2..48f243df48 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py @@ -49,27 +49,58 @@ def hflx_limiter_mo_stencil_03_min_max( @field_operator -def _hflx_limiter_mo_stencil_03( +def _hflx_limiter_mo_stencil_03a( z_mflx_anti_in: Field[[CellDim, KDim], float], z_mflx_anti_out: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], z_max: Field[[CellDim, KDim], float], z_min: Field[[CellDim, KDim], float], dbl_eps: float, + ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + return r_p, r_m +@field_operator +def _hflx_limiter_mo_stencil_03( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + # z_max: Field[[CellDim, KDim], float], + # z_min: Field[[CellDim, KDim], float], + dbl_eps: float, + +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + + z_max, z_min = _hflx_limiter_mo_stencil_03_min_max( z_tracer_max, z_tracer_min, beta_fct, r_beta_fct ) + + r_p, r_m =_hflx_limiter_mo_stencil_03a( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_max, + z_min, + dbl_eps, + ) + return r_p, r_m + @program def hflx_limiter_mo_stencil_03( z_tracer_max: Field[[CellDim, KDim], float], z_tracer_min: Field[[CellDim, KDim], float], beta_fct: float, r_beta_fct: float, - z_max: Field[[CellDim, KDim], float], - z_min: Field[[CellDim, KDim], float], + # z_max: Field[[CellDim, KDim], float], + # z_min: Field[[CellDim, KDim], float], z_mflx_anti_in: Field[[CellDim, KDim], float], z_mflx_anti_out: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], @@ -78,15 +109,16 @@ def hflx_limiter_mo_stencil_03( r_m: Field[[CellDim, KDim], float], ): - _hflx_limiter_mo_stencil_03_min_max( - z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_min, z_max) - ) _hflx_limiter_mo_stencil_03( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, - z_max, - z_min, + # z_max, + # z_min, dbl_eps, out=(r_p, r_m), - ) + ) \ No newline at end of file diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py index f696b358b3..ff05a66e3a 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -84,8 +84,8 @@ def test_hflx_diffusion_mo_stencil_03(): mesh = SimpleMesh() z_tracer_max = random_field(mesh, CellDim, KDim) z_tracer_min = random_field(mesh, CellDim, KDim) - z_max = zero_field(mesh, CellDim, KDim) - z_min = zero_field(mesh, CellDim, KDim) + # z_max = zero_field(mesh, CellDim, KDim) + # z_min = zero_field(mesh, CellDim, KDim) beta_fct = 0.4 r_beta_fct = 0.6 z_mflx_anti_in = random_field(mesh, CellDim, KDim) @@ -112,8 +112,8 @@ def test_hflx_diffusion_mo_stencil_03(): z_tracer_min, beta_fct, r_beta_fct, - z_max, - z_min, + # z_max, + # z_min, z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, From c71af473c18ee59d9225d63cc0515d6635259a21 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 Mar 2023 10:03:53 +0100 Subject: [PATCH 037/105] first draft for 5 --- .../advection/hflx_limiter_mo_stencil_05.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py new file mode 100644 index 0000000000..d5f342766d --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py @@ -0,0 +1,58 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import max_over, maximum, min_over, minimum, where + +from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, EdgeDim, E2C + + +@field_operator +def _hflx_limiter_mo_stencil_05( + z_anti: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + + z_signum = where( (z_anti>0.), 1., -1.) + + r_frac = 0.5*( (1. + z_signum)* + minimum(r_m(E2C[0]), + r_p(E2C[1]) ) + + (1.-z_signum)* + minimum( r_m(E2C[1]), + r_p(E2C[0]) ) ) + + p_mflx_tracer_h = z_mflx_low + minimum(1.,r_frac) * z_anti + + return p_mflx_tracer_h + + +@program +def hflx_limiter_mo_stencil_05( + z_anti: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_05( + z_anti, + z_mflx_low, + r_m, + r_p, + out=p_mflx_tracer_h + ) + From 539aa10ad6ff6fa4485f7484ec2d8f509d746d20 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 23 Mar 2023 18:01:03 +0100 Subject: [PATCH 038/105] fix merge --- advection/tests/test_vlimit_prbl_sm_stencil_01.py | 2 +- requirements.txt | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 9e53f46c1f..2765918b86 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -30,7 +30,7 @@ def v_limit_prbl_sm_stencil_01_numpy( z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:,:-1] + p_face[:,1:])) q_face_up, q_face_low = np.where( - abs(z_delta) < -1 * z_a6i, + np.abs(z_delta) < -1 * z_a6i, np.where( (p_cc < np.minimum(p_face[:,:-1], p_face[:,1:])), (p_cc, p_cc), diff --git a/requirements.txt b/requirements.txt index 6a508f229b..39fb3af434 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,8 +5,5 @@ ./pyutils ./testutils ./atm_dyn_iconam -<<<<<<< HEAD ./advection -======= ./liskov ->>>>>>> origin/main From 96b7222ba853e58a408531bb272edd1ba7eeb4e0 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 23 Mar 2023 18:38:10 +0100 Subject: [PATCH 039/105] cleanup --- .../src/icon4py/advection/hflx_limiter_mo_stencil_03.py | 6 ------ advection/tests/test_hflx_limiter_mo_stencil_03.py | 4 ---- 2 files changed, 10 deletions(-) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py index 48f243df48..a41ae7913c 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py @@ -74,8 +74,6 @@ def _hflx_limiter_mo_stencil_03( z_mflx_anti_in: Field[[CellDim, KDim], float], z_mflx_anti_out: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], - # z_max: Field[[CellDim, KDim], float], - # z_min: Field[[CellDim, KDim], float], dbl_eps: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: @@ -99,8 +97,6 @@ def hflx_limiter_mo_stencil_03( z_tracer_min: Field[[CellDim, KDim], float], beta_fct: float, r_beta_fct: float, - # z_max: Field[[CellDim, KDim], float], - # z_min: Field[[CellDim, KDim], float], z_mflx_anti_in: Field[[CellDim, KDim], float], z_mflx_anti_out: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], @@ -117,8 +113,6 @@ def hflx_limiter_mo_stencil_03( z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, - # z_max, - # z_min, dbl_eps, out=(r_p, r_m), ) \ No newline at end of file diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py index ff05a66e3a..d2901d96ca 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -84,8 +84,6 @@ def test_hflx_diffusion_mo_stencil_03(): mesh = SimpleMesh() z_tracer_max = random_field(mesh, CellDim, KDim) z_tracer_min = random_field(mesh, CellDim, KDim) - # z_max = zero_field(mesh, CellDim, KDim) - # z_min = zero_field(mesh, CellDim, KDim) beta_fct = 0.4 r_beta_fct = 0.6 z_mflx_anti_in = random_field(mesh, CellDim, KDim) @@ -112,8 +110,6 @@ def test_hflx_diffusion_mo_stencil_03(): z_tracer_min, beta_fct, r_beta_fct, - # z_max, - # z_min, z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, From 540b987b96d63b6f8c2bc2882a11baa60d81fddf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 3 Apr 2023 14:14:14 +0200 Subject: [PATCH 040/105] Fixed code generation bug in case there are no tolerance fields --- pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py index 19da5f1866..321dae71c2 100644 --- a/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py @@ -42,7 +42,11 @@ {%- for field in _this_node.out_fields -%} {{ field.renderer.render_ctype('c++') }} {{ field.renderer.render_pointer() }} {{ field.name }}_{{ suffix }}, {%- endfor -%} + {%- if _this_node.tol_fields -%} const int verticalStart, const int verticalEnd, const int horizontalStart, const int horizontalEnd, + {%- else -%} + const int verticalStart, const int verticalEnd, const int horizontalStart, const int horizontalEnd + {%- endif -%} {%- for field in _this_node.tol_fields -%} const double {{ field.name }}_rel_tol, const double {{ field.name }}_abs_tol From 6c3faacfb010e0b2425940c9e42856de5af2308f Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 12 Apr 2023 14:54:31 +0200 Subject: [PATCH 041/105] change from acc parallel to acc kernels --- liskov/README.md | 4 ++-- liskov/src/icon4py/liskov/codegen/f90.py | 4 ++-- liskov/tests/samples/fortran_samples.py | 4 ++-- liskov/tests/test_generation.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/liskov/README.md b/liskov/README.md index ffa4a5580f..b8ef8a913c 100644 --- a/liskov/README.md +++ b/liskov/README.md @@ -81,9 +81,9 @@ Together, the `START STENCIL` and `END STENCIL` directives result in the followi ```fortran #ifdef __DSL_VERIFY -!$ACC PARALLEL IF( i_am_accel_node .AND. acc_on ) DEFAULT(NONE) ASYNC(1) +!$ACC KERNELS IF( i_am_accel_node .AND. acc_on ) DEFAULT(NONE) ASYNC(1) vn_before(:, :, :) = vn(:, :, :) -!$ACC END PARALLEL +!$ACC END KERNELS ``` ```fortran diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index 5e8430d203..b8245344b3 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -322,11 +322,11 @@ class StartStencilStatementGenerator(TemplatedGenerator): """ #ifdef __DSL_VERIFY {% if _this_node.stencil_data.copies -%} - !$ACC PARALLEL IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) + !$ACC KERNELS IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) {%- for d in _this_node.copy_declarations %} {{ d.variable }}_before{{ d.lh_index }} = {{ d.association }}{{ d.rh_index }} {%- endfor %} - !$ACC END PARALLEL + !$ACC END KERNELS {%- endif -%} {%- if _this_node.profile %} diff --git a/liskov/tests/samples/fortran_samples.py b/liskov/tests/samples/fortran_samples.py index 858b3a2089..5a9a5f03c2 100644 --- a/liskov/tests/samples/fortran_samples.py +++ b/liskov/tests/samples/fortran_samples.py @@ -48,9 +48,9 @@ CALL get_indices_e(p_patch, jb, i_startblk, i_endblk, & i_startidx, i_endidx, start_bdydiff_e, grf_bdywidth_e) - !$ACC PARALLEL IF( i_am_accel_node .AND. acc_on ) DEFAULT(NONE) ASYNC(1) + !$ACC KERNELS IF( i_am_accel_node .AND. acc_on ) DEFAULT(NONE) ASYNC(1) vn_before(:,:,:) = p_nh_prog%vn(:,:,:) - !$ACC END PARALLEL + !$ACC END KERNELS !$ACC PARALLEL LOOP DEFAULT(NONE) GANG VECTOR COLLAPSE(2) ASYNC(1) IF( i_am_accel_node .AND. acc_on ) DO jk = 1, nlev diff --git a/liskov/tests/test_generation.py b/liskov/tests/test_generation.py index a83f66b8e5..f41a33e4ac 100644 --- a/liskov/tests/test_generation.py +++ b/liskov/tests/test_generation.py @@ -141,14 +141,14 @@ def expected_declare_source(): def expected_start_stencil_source(): return """ #ifdef __DSL_VERIFY - !$ACC PARALLEL IF( i_am_accel_node ) DEFAULT(NONE) ASYNC(1) + !$ACC KERNELS IF( i_am_accel_node ) DEFAULT(NONE) ASYNC(1) out1_before(:, :) = out1(:, :, 1) out2_before(:, :, :) = p_nh%prog(nnew)%out2(:, :, :) out3_before(:, :) = p_nh%prog(nnew)%w(:, :, jb) out4_before(:, :, :) = p_nh%prog(nnew)%w(:, :, :, 2) out5_before(:, :, :) = p_nh%prog(nnew)%w(:, :, :, ntnd) out6_before(:, :, :) = p_nh%prog(nnew)%w(:, :, :, ntnd) - !$ACC END PARALLEL + !$ACC END KERNELS call nvtxStartRange("stencil1")""" From 548591fa93d21a259b4f3ea06df9c91b3e8a2b0f Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 4 May 2023 15:13:09 +0200 Subject: [PATCH 042/105] first draft --- liskov/src/icon4py/liskov/codegen/f90.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index 760f4e7539..d7463c69b8 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -164,6 +164,10 @@ class EndStencilStatementGenerator(TemplatedGenerator): {{ output_fields }} {{ tolerance_fields }} {{ bounds_fields }} + + #ifdef __DSL_VERIFY + !$ACC END DATA + #endif """ ) @@ -321,6 +325,13 @@ class StartStencilStatementGenerator(TemplatedGenerator): StartStencilStatement = as_jinja( """ #ifdef __DSL_VERIFY + + !$ACC DATA CREATE( & + {%- for d in _this_node.copy_declarations %} + !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} + {%- endfor %} + !$ACC IF ( i_am_accel_node ) + {% if _this_node.stencil_data.copies -%} !$ACC KERNELS IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) {%- for d in _this_node.copy_declarations %} From c230cb567f45f680e21fb809628b4539aa88ab2c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 4 May 2023 16:16:57 +0200 Subject: [PATCH 043/105] first cleanup startcreate --- liskov/src/icon4py/liskov/codegen/f90.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index d7463c69b8..df22bb2438 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -385,14 +385,10 @@ class StartCreateStatementGenerator(TemplatedGenerator): !$ACC DATA CREATE( & {%- if _this_node.extra_fields -%} {%- for name in extra_fields %} - !$ACC {{ name }} {%- if not loop.last -%}, & {% else %}, & {%- endif -%} + !$ACC {{ name }} {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} {%- endfor %} - {%- endif -%} - {%- for name in out_field_names %} - !$ACC {{ name }}_before {%- if not loop.last -%}, & {% else %} & {%- endif -%} - {%- endfor %} - !$ACC ), & - !$ACC IF ( i_am_accel_node ) + {%- endif %} + !$ACC IF ( i_am_accel_node ) """ ) From bbae0c0f25376a74b78a2e2e7133de6800093a9d Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 4 May 2023 16:26:13 +0200 Subject: [PATCH 044/105] update validate --- liskov/src/icon4py/liskov/parsing/validation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/liskov/src/icon4py/liskov/parsing/validation.py b/liskov/src/icon4py/liskov/parsing/validation.py index 616803e709..a489262b63 100644 --- a/liskov/src/icon4py/liskov/parsing/validation.py +++ b/liskov/src/icon4py/liskov/parsing/validation.py @@ -117,6 +117,8 @@ def _validate_directive_uniqueness( repeated = remove_directive_types( [d for d in directives if directives.count(d) > 1], [ + ts.StartCreate, + ts.EndCreate, ts.StartStencil, ts.EndStencil, ts.EndIf, @@ -138,8 +140,6 @@ def _validate_required_directives( expected = [ ts.Declare, ts.Imports, - ts.StartCreate, - ts.EndCreate, ts.StartStencil, ts.EndStencil, ] From 07c6bf74ec512622e4eceba92da7b1ebe265b295 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 4 May 2023 16:50:53 +0200 Subject: [PATCH 045/105] make START CREATE with Arguements --- liskov/src/icon4py/liskov/parsing/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/liskov/src/icon4py/liskov/parsing/types.py b/liskov/src/icon4py/liskov/parsing/types.py index 9fb5ee29da..c64c040272 100644 --- a/liskov/src/icon4py/liskov/parsing/types.py +++ b/liskov/src/icon4py/liskov/parsing/types.py @@ -128,7 +128,7 @@ class Imports(WithoutArguments): pattern = "IMPORTS" -class StartCreate(WithOptionalArguments): +class StartCreate(WithArguments): pattern = "START CREATE" From a8b21a972f207201dba555083a1fa7c52bd67918 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 9 May 2023 17:22:58 +0200 Subject: [PATCH 046/105] kind of working --- liskov/src/icon4py/liskov/codegen/f90.py | 15 ---------- liskov/src/icon4py/liskov/codegen/generate.py | 30 +++++++++---------- .../src/icon4py/liskov/codegen/interface.py | 4 +-- .../src/icon4py/liskov/parsing/deserialise.py | 29 +++++++++++------- 4 files changed, 36 insertions(+), 42 deletions(-) diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index df22bb2438..26bfba6b4b 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -362,22 +362,7 @@ class ImportsStatementGenerator(TemplatedGenerator): class StartCreateStatement(eve.Node): - stencils: list[StartStencilData] extra_fields: Optional[list[str]] - out_field_names: list[str] = eve.datamodels.field(init=False) - - def __post_init__(self) -> None: # type: ignore - self.out_field_names = sorted( - set( - [ - field.variable - for stencil in self.stencils - for field in stencil.fields - if field.out - ] - ) - ) - class StartCreateStatementGenerator(TemplatedGenerator): StartCreateStatement = as_jinja( diff --git a/liskov/src/icon4py/liskov/codegen/generate.py b/liskov/src/icon4py/liskov/codegen/generate.py index 26db2e626f..b1e069d37d 100644 --- a/liskov/src/icon4py/liskov/codegen/generate.py +++ b/liskov/src/icon4py/liskov/codegen/generate.py @@ -220,22 +220,22 @@ def _generate_imports(self) -> None: def _generate_create(self) -> None: """Generate f90 code for OpenACC DATA CREATE statements.""" - logger.info("Generating DATA CREATE statement.") - self._generate( - StartCreateStatement, - StartCreateStatementGenerator, - self.directives.StartCreate.startln, - self.directives.StartCreate.endln, - stencils=self.directives.StartStencil, - extra_fields=self.directives.StartCreate.extra_fields, - ) + if self.directives.StartCreate != UnusedDirective: + logger.info("Generating DATA CREATE statement.") + self._generate( + StartCreateStatement, + StartCreateStatementGenerator, + self.directives.StartCreate.startln, + self.directives.StartCreate.endln, + extra_fields=self.directives.StartCreate.extra_fields, + ) - self._generate( - EndCreateStatement, - EndCreateStatementGenerator, - self.directives.EndCreate.startln, - self.directives.EndCreate.endln, - ) + self._generate( + EndCreateStatement, + EndCreateStatementGenerator, + self.directives.EndCreate.startln, + self.directives.EndCreate.endln, + ) def _generate_endif(self) -> None: """Generate f90 code for endif statements.""" diff --git a/liskov/src/icon4py/liskov/codegen/interface.py b/liskov/src/icon4py/liskov/codegen/interface.py index dfa879f0cd..94e078040b 100644 --- a/liskov/src/icon4py/liskov/codegen/interface.py +++ b/liskov/src/icon4py/liskov/codegen/interface.py @@ -109,8 +109,8 @@ class DeserialisedDirectives: EndStencil: Sequence[EndStencilData] Declare: Sequence[DeclareData] Imports: ImportsData - StartCreate: StartCreateData - EndCreate: EndCreateData + StartCreate: Sequence[StartCreateData] | UnusedDirective + EndCreate: Sequence[EndCreateData] | UnusedDirective EndIf: Sequence[EndIfData] | UnusedDirective StartProfile: Sequence[StartProfileData] | UnusedDirective EndProfile: Sequence[EndProfileData] | UnusedDirective diff --git a/liskov/src/icon4py/liskov/parsing/deserialise.py b/liskov/src/icon4py/liskov/parsing/deserialise.py index fa88ae048c..6455d8fb35 100644 --- a/liskov/src/icon4py/liskov/parsing/deserialise.py +++ b/liskov/src/icon4py/liskov/parsing/deserialise.py @@ -122,7 +122,7 @@ def __call__(self, parsed: ts.ParsedDict) -> CodeGenInput: @dataclass -class EndCreateDataFactory(RequiredSingleUseDataFactory): +class EndCreateDataFactory(OptionalMultiUseDataFactory): directive_cls: Type[ts.ParsedDirective] = ts.EndCreate dtype: Type[EndCreateData] = EndCreateData @@ -150,18 +150,27 @@ class StartCreateDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = ts.StartCreate dtype: Type[StartCreateData] = StartCreateData - def __call__(self, parsed: ts.ParsedDict) -> StartCreateData: - directive = extract_directive(parsed["directives"], self.directive_cls)[0] + def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData]: - named_args = parsed["content"]["StartCreate"][0] + deserialised = [] + extracted = extract_directive(parsed["directives"], self.directive_cls) - extra_fields = None - if named_args: - extra_fields = named_args["extra_fields"].split(",") + if len(extracted) < 1: + return UnusedDirective - return self.dtype( - startln=directive.startln, endln=directive.endln, extra_fields=extra_fields - ) + for i, directive in enumerate(extracted): + + named_args = parsed["content"]["StartCreate"][i] + + extra_fields = None + if named_args: + extra_fields = named_args["extra_fields"].split(",") + + deserialised.append( + self.dtype( startln=directive.startln, endln=directive.endln, extra_fields=extra_fields ) + ) + + return deserialised @dataclass From 3400755e6dc007b9b6fd66a6bf644416e408991d Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 9 May 2023 17:40:32 +0200 Subject: [PATCH 047/105] fix generate --- liskov/src/icon4py/liskov/codegen/generate.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/liskov/src/icon4py/liskov/codegen/generate.py b/liskov/src/icon4py/liskov/codegen/generate.py index b1e069d37d..fcf46d0a3d 100644 --- a/liskov/src/icon4py/liskov/codegen/generate.py +++ b/liskov/src/icon4py/liskov/codegen/generate.py @@ -221,21 +221,24 @@ def _generate_imports(self) -> None: def _generate_create(self) -> None: """Generate f90 code for OpenACC DATA CREATE statements.""" if self.directives.StartCreate != UnusedDirective: - logger.info("Generating DATA CREATE statement.") - self._generate( - StartCreateStatement, - StartCreateStatementGenerator, - self.directives.StartCreate.startln, - self.directives.StartCreate.endln, - extra_fields=self.directives.StartCreate.extra_fields, - ) + for startcreate in self.directives.StartCreate: # type: ignore + logger.info("Generating DATA CREATE statement.") + self._generate( + StartCreateStatement, + StartCreateStatementGenerator, + startcreate.startln, + startcreate.endln, + extra_fields=startcreate.extra_fields, + ) - self._generate( - EndCreateStatement, - EndCreateStatementGenerator, - self.directives.EndCreate.startln, - self.directives.EndCreate.endln, - ) + if self.directives.EndCreate != UnusedDirective: + for endcreate in self.directives.EndCreate: # type: ignore + self._generate( + EndCreateStatement, + EndCreateStatementGenerator, + endcreate.startln, + endcreate.endln, + ) def _generate_endif(self) -> None: """Generate f90 code for endif statements.""" From 0db9940a1257693c3081dc6a48926caf7cb74e11 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 15 May 2023 14:10:12 +0200 Subject: [PATCH 048/105] add accnodata --- liskov/src/icon4py/liskov/codegen/f90.py | 3 +++ liskov/src/icon4py/liskov/codegen/generate.py | 1 + liskov/src/icon4py/liskov/codegen/interface.py | 1 + liskov/src/icon4py/liskov/parsing/deserialise.py | 2 ++ liskov/tests/test_validation.py | 3 --- 5 files changed, 7 insertions(+), 3 deletions(-) diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index 26bfba6b4b..e4afc05638 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -134,6 +134,7 @@ class EndStencilStatement(eve.Node): profile: bool noendif: Optional[bool] noprofile: Optional[bool] + noaccenddata: Optional[bool] name: str = eve.datamodels.field(init=False) input_fields: InputFields = eve.datamodels.field(init=False) @@ -165,9 +166,11 @@ class EndStencilStatementGenerator(TemplatedGenerator): {{ tolerance_fields }} {{ bounds_fields }} + {%- if not _this_node.noaccenddata %} #ifdef __DSL_VERIFY !$ACC END DATA #endif + {%- endif %} """ ) diff --git a/liskov/src/icon4py/liskov/codegen/generate.py b/liskov/src/icon4py/liskov/codegen/generate.py index fcf46d0a3d..83fa0c1b41 100644 --- a/liskov/src/icon4py/liskov/codegen/generate.py +++ b/liskov/src/icon4py/liskov/codegen/generate.py @@ -205,6 +205,7 @@ def _generate_end_stencil(self) -> None: profile=self.profile, noendif=self.directives.EndStencil[i].noendif, noprofile=self.directives.EndStencil[i].noprofile, + noaccenddata=self.directives.EndStencil[i].noaccenddata ) def _generate_imports(self) -> None: diff --git a/liskov/src/icon4py/liskov/codegen/interface.py b/liskov/src/icon4py/liskov/codegen/interface.py index 94e078040b..d93160ec54 100644 --- a/liskov/src/icon4py/liskov/codegen/interface.py +++ b/liskov/src/icon4py/liskov/codegen/interface.py @@ -96,6 +96,7 @@ class EndStencilData(CodeGenInput): name: str noendif: Optional[bool] noprofile: Optional[bool] + noaccenddata: Optional[bool] @dataclass diff --git a/liskov/src/icon4py/liskov/parsing/deserialise.py b/liskov/src/icon4py/liskov/parsing/deserialise.py index 6455d8fb35..8bb01ac928 100644 --- a/liskov/src/icon4py/liskov/parsing/deserialise.py +++ b/liskov/src/icon4py/liskov/parsing/deserialise.py @@ -237,6 +237,7 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndStencilData]: stencil_name = _extract_stencil_name(named_args, directive) noendif = _extract_boolean_kwarg(directive, named_args, "noendif") noprofile = _extract_boolean_kwarg(directive, named_args, "noprofile") + noaccenddata = _extract_boolean_kwarg(directive, named_args, "noaccenddata") deserialised.append( self.dtype( name=stencil_name, @@ -244,6 +245,7 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndStencilData]: endln=directive.endln, noendif=noendif, noprofile=noprofile, + noaccenddata=noaccenddata, ) ) return deserialised diff --git a/liskov/tests/test_validation.py b/liskov/tests/test_validation.py index 03bd6e5b53..885c70d3bb 100644 --- a/liskov/tests/test_validation.py +++ b/liskov/tests/test_validation.py @@ -63,7 +63,6 @@ def test_directive_semantics_validation_unbalanced_stencil_directives( [Declare("!$DSL DECLARE(name=foo; bar)", 0, 0)], [Imports("!$DSL IMPORTS(foo)", 0, 0)], [Imports("!$DSL IMPORTS())", 0, 0)], - [StartCreate("!$DSL START CREATE(;)", 0, 0)], ), ) def test_directive_syntax_validator(directive): @@ -76,7 +75,6 @@ def test_directive_syntax_validator(directive): "directive", [ "!$DSL IMPORTS()", - "!$DSL START CREATE()", ], ) def test_directive_semantics_validation_repeated_directives( @@ -114,7 +112,6 @@ def test_directive_semantics_validation_repeated_stencil(make_f90_tmpfile, direc "directive", [ """!$DSL IMPORTS()""", - """!$DSL START CREATE()""", """!$DSL END STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary; noprofile=True)""", ], ) From 928ac5e22354cdd93588399aad21a1a3729f4614 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 17 May 2023 15:48:33 +0200 Subject: [PATCH 049/105] fix before fields --- liskov/src/icon4py/liskov/codegen/f90.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index e4afc05638..33f72bae2d 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -167,9 +167,7 @@ class EndStencilStatementGenerator(TemplatedGenerator): {{ bounds_fields }} {%- if not _this_node.noaccenddata %} - #ifdef __DSL_VERIFY !$ACC END DATA - #endif {%- endif %} """ ) @@ -327,7 +325,6 @@ def render_index(n: int) -> str: class StartStencilStatementGenerator(TemplatedGenerator): StartStencilStatement = as_jinja( """ - #ifdef __DSL_VERIFY !$ACC DATA CREATE( & {%- for d in _this_node.copy_declarations %} @@ -335,6 +332,7 @@ class StartStencilStatementGenerator(TemplatedGenerator): {%- endfor %} !$ACC IF ( i_am_accel_node ) + #ifdef __DSL_VERIFY {% if _this_node.stencil_data.copies -%} !$ACC KERNELS IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) {%- for d in _this_node.copy_declarations %} From b6966bfb8922654803d9caacb47c68fc2bb4c538 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 17 May 2023 17:50:31 +0200 Subject: [PATCH 050/105] fix first testss --- liskov/tests/test_generation.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/liskov/tests/test_generation.py b/liskov/tests/test_generation.py index f41a33e4ac..1903c60f01 100644 --- a/liskov/tests/test_generation.py +++ b/liskov/tests/test_generation.py @@ -71,7 +71,7 @@ def serialised_directives(): copies=True, ) end_stencil_data = EndStencilData( - name="stencil1", startln=3, endln=4, noendif=False, noprofile=False + name="stencil1", startln=3, endln=4, noendif=False, noprofile=False, noaccenddata=False ) declare_data = DeclareData( startln=5, @@ -95,8 +95,8 @@ def serialised_directives(): EndStencil=[end_stencil_data], Declare=[declare_data], Imports=imports_data, - StartCreate=start_create_data, - EndCreate=end_create_data, + StartCreate=[start_create_data], + EndCreate=[end_create_data], EndIf=[endif_data], StartProfile=[start_profile_data], EndProfile=[end_profile_data], @@ -109,15 +109,8 @@ def expected_start_create_source(): return """ !$ACC DATA CREATE( & !$ACC foo, & -!$ACC bar, & -!$ACC out1_before, & -!$ACC out2_before, & -!$ACC out3_before, & -!$ACC out4_before, & -!$ACC out5_before, & -!$ACC out6_before & -!$ACC ), & -!$ACC IF ( i_am_accel_node )""" +!$ACC bar ) & +!$ACC IF ( i_am_accel_node )""" @pytest.fixture From 27731489fd303bb1a1fbd2d7719ec1d25a004561 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 15:29:59 +0200 Subject: [PATCH 051/105] fix test_generation --- liskov/tests/test_generation.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/liskov/tests/test_generation.py b/liskov/tests/test_generation.py index 1903c60f01..6b4a6d8926 100644 --- a/liskov/tests/test_generation.py +++ b/liskov/tests/test_generation.py @@ -133,6 +133,15 @@ def expected_declare_source(): @pytest.fixture def expected_start_stencil_source(): return """ + !$ACC DATA CREATE( & + !$ACC out1_before, & + !$ACC out2_before, & + !$ACC out3_before, & + !$ACC out4_before, & + !$ACC out5_before, & + !$ACC out6_before ) & + !$ACC IF ( i_am_accel_node ) + #ifdef __DSL_VERIFY !$ACC KERNELS IF( i_am_accel_node ) DEFAULT(NONE) ASYNC(1) out1_before(:, :) = out1(:, :, 1) @@ -170,7 +179,9 @@ def expected_end_stencil_source(): vertical_lower=-1, & vertical_upper=-10, & horizontal_lower=1, & - horizontal_upper=10)""" + horizontal_upper=10) + + !$ACC END DATA""" @pytest.fixture From f02de6222dfb7e66421a3c14e7843015737f2317 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 16:14:27 +0200 Subject: [PATCH 052/105] fix test --- liskov/tests/samples/fortran_samples.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/liskov/tests/samples/fortran_samples.py b/liskov/tests/samples/fortran_samples.py index 5a9a5f03c2..e16d6fc313 100644 --- a/liskov/tests/samples/fortran_samples.py +++ b/liskov/tests/samples/fortran_samples.py @@ -29,8 +29,6 @@ SINGLE_STENCIL = """\ !$DSL IMPORTS() - !$DSL START CREATE() - !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; suffix=dsl) !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; a=nproma,p_patch%nlev,p_patch%nblks_e; & @@ -66,14 +64,11 @@ !$ACC END PARALLEL LOOP !$DSL END PROFILE() !$DSL END STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary; noprofile=True) - !$DSL END CREATE() """ MULTIPLE_STENCILS = """\ !$DSL IMPORTS() - !$DSL START CREATE() - !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; z_rho_e=nproma,p_patch%nlev,p_patch%nblks_c; & !$DSL z_theta_v_e=nproma,p_patch%nlev,p_patch%nblks_c; z_nabla2_c=nproma,p_patch%nlev,p_patch%nblks_c; & !$DSL z_rth_pr_1=nproma,p_patch%nlev,p_patch%nblks_c; z_rth_pr_2=nproma,p_patch%nlev,p_patch%nblks_c; & @@ -155,7 +150,6 @@ !$ACC END PARALLEL LOOP !$DSL ENDIF() !$DSL END STENCIL(name=calculate_nabla2_for_w; noendif=true) - !$DSL END CREATE() """ DIRECTIVES_SAMPLE = """\ From 19debe573d75d457077f190792a447f1f8b6825b Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 16:53:24 +0200 Subject: [PATCH 053/105] fix parser --- liskov/tests/test_parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/liskov/tests/test_parser.py b/liskov/tests/test_parser.py index 8101988c51..a8423abf60 100644 --- a/liskov/tests/test_parser.py +++ b/liskov/tests/test_parser.py @@ -79,7 +79,7 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont @mark.parametrize( "stencil, num_directives, num_content", - [(SINGLE_STENCIL, 9, 8), (MULTIPLE_STENCILS, 11, 7)], + [(SINGLE_STENCIL, 7, 6), (MULTIPLE_STENCILS, 9, 5)], ) def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): fpath = make_f90_tmpfile(content=stencil) From 60bb67107251e74fc058b89145906458e307e685 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 17:05:50 +0200 Subject: [PATCH 054/105] fix test --- liskov/tests/test_deserialiser.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/liskov/tests/test_deserialiser.py b/liskov/tests/test_deserialiser.py index f32c4e2c04..f5e3c4aaf5 100644 --- a/liskov/tests/test_deserialiser.py +++ b/liskov/tests/test_deserialiser.py @@ -195,6 +195,8 @@ def test_data_factories_with_args(factory, target, mock_data): def test_start_create_factory(mock_data, extra_fields): factory = StartCreateDataFactory() result = factory(mock_data) + if type(result) == list: + result = result[0] assert isinstance(result, StartCreateData) assert result.extra_fields == extra_fields From 0552f7445a154774210a50bd09b5ced8e4b27207 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 17:13:38 +0200 Subject: [PATCH 055/105] fix test --- liskov/tests/samples/fortran_samples.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/liskov/tests/samples/fortran_samples.py b/liskov/tests/samples/fortran_samples.py index e16d6fc313..bd2d262106 100644 --- a/liskov/tests/samples/fortran_samples.py +++ b/liskov/tests/samples/fortran_samples.py @@ -175,8 +175,6 @@ CONSECUTIVE_STENCIL = """\ !$DSL IMPORTS() - !$DSL START CREATE() - !$DSL DECLARE(z_q=nproma,p_patch%nlev; z_alpha=nproma,p_patch%nlev) !$DSL START STENCIL(name=mo_solve_nonhydro_stencil_45; z_alpha=z_alpha(:,:); vertical_lower=nlevp1; & @@ -202,15 +200,12 @@ !$DSL END STENCIL(name=mo_solve_nonhydro_stencil_45; noendif=true; noprofile=true) !$DSL END STENCIL(name=mo_solve_nonhydro_stencil_45_b; noendif=true; noprofile=true) - !$DSL END CREATE() """ FREE_FORM_STENCIL = """\ !$DSL IMPORTS() - !$DSL START CREATE() - !$DSL DECLARE(z_q=nproma,p_patch%nlev; z_alpha=nproma,p_patch%nlev) !$DSL INSERT(some custom fields go here) @@ -234,5 +229,4 @@ !$DSL END STENCIL(name=mo_solve_nonhydro_stencil_45) - !$DSL END CREATE() """ From 2c6292e862744566525a74e4b291331d73699b50 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 17:41:59 +0200 Subject: [PATCH 056/105] update README --- liskov/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/liskov/README.md b/liskov/README.md index b8ef8a913c..23b32420f7 100644 --- a/liskov/README.md +++ b/liskov/README.md @@ -30,7 +30,7 @@ This directive generates the necessary `USE` statements to import the Fortran to #### `!$DSL START CREATE()` -This directive generates an OpenACC `DATA CREATE` statement for all output fields used in each DSL (icon4py) stencil. The directive also takes an **optional** keyword argument to specify extra fields to include in the `DATA CREATE` statement called `extra_fields`. Here you can specify a comma-separated list of strings which should be added to the `DATA CREATE` statement as follows `extra_fields=foo,bar`. +This directive generates an OpenACC `DATA CREATE` statement. The directive takes an **optional** keyword argument to specify extra fields to include in the `DATA CREATE` statement called `extra_fields`. Here you can specify a comma-separated list of strings which should be added to the `DATA CREATE` statement as follows `extra_fields=foo,bar`. #### `!$DSL END CREATE()` From 7a560a74cd148333589abc99cf876b43030853d4 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 23 May 2023 17:42:54 +0200 Subject: [PATCH 057/105] impprove README --- liskov/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/liskov/README.md b/liskov/README.md index 23b32420f7..bf02b1e3ef 100644 --- a/liskov/README.md +++ b/liskov/README.md @@ -30,7 +30,7 @@ This directive generates the necessary `USE` statements to import the Fortran to #### `!$DSL START CREATE()` -This directive generates an OpenACC `DATA CREATE` statement. The directive takes an **optional** keyword argument to specify extra fields to include in the `DATA CREATE` statement called `extra_fields`. Here you can specify a comma-separated list of strings which should be added to the `DATA CREATE` statement as follows `extra_fields=foo,bar`. +This directive generates an OpenACC `DATA CREATE` statement. The directive requires an **optional** keyword argument to specify extra fields to include in the `DATA CREATE` statement called `extra_fields`. Here you can specify a comma-separated list of strings which should be added to the `DATA CREATE` statement as follows `extra_fields=foo,bar`. #### `!$DSL END CREATE()` From 366a0bbe856c065ce3987ef1f00c198786c8e249 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 25 May 2023 09:12:06 +0200 Subject: [PATCH 058/105] fix style --- .../advection/btraj_dreg_stencil_01.py | 9 +- .../advection/btraj_dreg_stencil_02.py | 16 +- .../advection/btraj_dreg_stencil_03.py | 95 +++- .../divide_flux_area_list_stencil_01.py | 462 +++++++++++++++--- .../divide_flux_area_list_stencil_02.py | 192 +++++++- .../advection/face_val_ppm_stencil_01.py | 30 +- .../advection/face_val_ppm_stencil_02.py | 23 +- .../advection/hflux_ffsl_hybrid_stencil_02.py | 5 +- .../advection/hflx_limiter_mo_stencil_01a.py | 11 +- .../advection/hflx_limiter_mo_stencil_01b.py | 46 +- .../advection/hflx_limiter_mo_stencil_02.py | 8 +- .../advection/hflx_limiter_mo_stencil_03.py | 12 +- .../advection/hflx_limiter_mo_stencil_05.py | 27 +- .../prep_gauss_quadrature_c_list_stencil.py | 362 +++++++++++--- .../prep_gauss_quadrature_c_stencil.py | 316 ++++++++++-- .../advection/recon_lsq_cell_c_stencil.py | 312 ++++++++++-- .../advection/recon_lsq_cell_c_svd_stencil.py | 181 ++++++- .../upwind_hflux_miura_stencil_02.py | 6 +- .../tests/test_face_val_ppm_stencil_01.py | 61 ++- .../tests/test_face_val_ppm_stencil_02.py | 26 +- .../tests/test_face_val_ppm_stencil_05.py | 18 +- .../tests/test_step_advection_stencil_01.py | 3 +- .../tests/test_step_advection_stencil_02.py | 4 +- advection/tests/test_vert_adv_stencil_01.py | 5 +- .../tests/test_vlimit_prbl_sm_stencil_01.py | 14 +- liskov/src/icon4py/liskov/codegen/f90.py | 1 + liskov/src/icon4py/liskov/codegen/generate.py | 34 +- .../src/icon4py/liskov/parsing/deserialise.py | 6 +- liskov/tests/test_generation.py | 7 +- liskov/tests/test_validation.py | 7 +- 30 files changed, 1866 insertions(+), 433 deletions(-) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_01.py b/advection/src/icon4py/advection/btraj_dreg_stencil_01.py index 541057a296..472c90f1f4 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_01.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_01.py @@ -11,10 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.ffront.fbuiltins import Field, int32, where, broadcast from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, where from icon4py.common.dimension import EdgeDim, KDim + + @field_operator def _btraj_dreg_stencil_01( lcounterclock: bool, @@ -23,9 +25,9 @@ def _btraj_dreg_stencil_01( ) -> Field[[EdgeDim, KDim], bool]: tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim)) - lvn_sys_pos_true = where(p_vn*tangent_orientation >= 0.0, True, False) + lvn_sys_pos_true = where(p_vn * tangent_orientation >= 0.0, True, False) mask_lcounterclock = broadcast(lcounterclock, (EdgeDim, KDim)) - lvn_sys_pos = where(mask_lcounterclock, lvn_sys_pos_true, False) + lvn_sys_pos = where(mask_lcounterclock, lvn_sys_pos_true, False) return lvn_sys_pos @@ -37,4 +39,3 @@ def btraj_dreg_stencil_01( lvn_sys_pos: Field[[EdgeDim, KDim], bool], ): _btraj_dreg_stencil_01(lcounterclock, p_vn, tangent_orientation, out=lvn_sys_pos) - diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py index 1fa5656dad..0b2a2f3630 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.ffront.fbuiltins import Field, where, sqrt, broadcast, int32 from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, sqrt, where from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim @@ -28,8 +28,10 @@ def _btraj_dreg_stencil_02( lvn_pos = where(p_vn >= 0.0, True, False) traj_length = sqrt(p_vn**2 + p_vt**2) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) - famask = where(traj_length > 1.25*broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0)) - + famask = where( + traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) + ) + return famask @@ -40,11 +42,5 @@ def btraj_dreg_stencil_02( edge_cell_length: Field[[ECDim], float], famask: Field[[EdgeDim, KDim], int32], p_dt: float, - ): - _btraj_dreg_stencil_02( - p_vn, - p_vt, - edge_cell_length, - p_dt, - out=famask) + _btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=famask) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py index 269ba06f8d..a43d1670f4 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.ffront.fbuiltins import Field, int32, where, broadcast from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim @@ -49,8 +49,8 @@ def _btraj_dreg_stencil_03( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: -# lvn_sys_pos = where(lvn_sys_pos_int == int32(1), broadcast(True, (EdgeDim, KDim)), broadcast(False, (EdgeDim, KDim))) -# lvn_sys_pos = where(lvn_sys_pos_int == int32(1), True, False) + # this is comented out: lvn_sys_pos = where(lvn_sys_pos_int == int32(1), broadcast(True, (EdgeDim, KDim)), broadcast(False, (EdgeDim, KDim))) + # this is comented out: lvn_sys_pos = where(lvn_sys_pos_int == int32(1), True, False) lvn_pos = where(p_vn >= 0.0, True, False) @@ -67,29 +67,67 @@ def _btraj_dreg_stencil_03( pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y - pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x - pos_dreg_vert_c_2_y = where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = ( + where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_2_y = ( + where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + ) pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y - pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x - pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - - pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) - pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) + pos_dreg_vert_c_4_x = ( + where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_4_y = ( + where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + ) + + pn_cell_1 = where( + lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1]) + ) + pn_cell_2 = where( + lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1]) + ) dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) - p_coords_dreg_v_1_lon = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 - p_coords_dreg_v_2_lon = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 - p_coords_dreg_v_3_lon = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 - p_coords_dreg_v_4_lon = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - p_coords_dreg_v_1_lat = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 - p_coords_dreg_v_2_lat = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 - p_coords_dreg_v_3_lat = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 - p_coords_dreg_v_4_lat = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 - - - return p_cell_idx, p_cell_blk, p_coords_dreg_v_1_lon, p_coords_dreg_v_2_lon, p_coords_dreg_v_3_lon, p_coords_dreg_v_4_lon, p_coords_dreg_v_1_lat, p_coords_dreg_v_2_lat, p_coords_dreg_v_3_lat, p_coords_dreg_v_4_lat + p_coords_dreg_v_1_lon = ( + pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + ) + p_coords_dreg_v_2_lon = ( + pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + ) + p_coords_dreg_v_3_lon = ( + pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + ) + p_coords_dreg_v_4_lon = ( + pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + ) + p_coords_dreg_v_1_lat = ( + pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + ) + p_coords_dreg_v_2_lat = ( + pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + ) + p_coords_dreg_v_3_lat = ( + pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + ) + p_coords_dreg_v_4_lat = ( + pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + ) + + return ( + p_cell_idx, + p_cell_blk, + p_coords_dreg_v_1_lon, + p_coords_dreg_v_2_lon, + p_coords_dreg_v_3_lon, + p_coords_dreg_v_4_lon, + p_coords_dreg_v_1_lat, + p_coords_dreg_v_2_lat, + p_coords_dreg_v_3_lat, + p_coords_dreg_v_4_lat, + ) @program @@ -122,7 +160,6 @@ def btraj_dreg_stencil_03( p_coords_dreg_v_2_lat: Field[[EdgeDim, KDim], float], p_coords_dreg_v_3_lat: Field[[EdgeDim, KDim], float], p_coords_dreg_v_4_lat: Field[[EdgeDim, KDim], float], - ): _btraj_dreg_stencil_03( p_vn, @@ -143,4 +180,16 @@ def btraj_dreg_stencil_03( dual_normal_cell_y, lvn_sys_pos, p_dt, - out=(p_cell_idx, p_cell_blk, p_coords_dreg_v_1_lon, p_coords_dreg_v_2_lon, p_coords_dreg_v_3_lon, p_coords_dreg_v_4_lon, p_coords_dreg_v_1_lat, p_coords_dreg_v_2_lat, p_coords_dreg_v_3_lat, p_coords_dreg_v_4_lat)) + out=( + p_cell_idx, + p_cell_blk, + p_coords_dreg_v_1_lon, + p_coords_dreg_v_2_lon, + p_coords_dreg_v_3_lon, + p_coords_dreg_v_4_lon, + p_coords_dreg_v_1_lat, + p_coords_dreg_v_2_lat, + p_coords_dreg_v_3_lat, + p_coords_dreg_v_4_lat, + ), + ) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py index 9f998ab353..98a5ecd63e 100644 --- a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py @@ -11,12 +11,17 @@ # # 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, broadcast, int32, abs -from icon4py.common.dimension import E2EC, ECDim, CellDim, EdgeDim, KDim import sys + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim + + sys.setrecursionlimit(5500) + # FUNCTIONS # Checking turn when travelling along three points, used to check whether lines inters. @field_operator @@ -39,9 +44,10 @@ def ccw( dy1dx2 = dy1 * dx2 lccw = where(dx1dy2 > dy1dx2, True, False) - ccw_out = where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + ccw_out = where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise return ccw_out + # Checks whether two lines intersect @field_operator def lintersect( @@ -55,12 +61,41 @@ def lintersect( line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], bool]: - intersect1 = ccw(line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p1_lon, line2_p1_lat) * ccw(line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p2_lon, line2_p2_lat) - intersect2 = ccw(line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, line1_p1_lon, line1_p1_lat) * ccw(line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, line1_p2_lon, line1_p2_lat) + intersect1 = ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) lintersect_out = where((intersect1 + intersect2) == -2, True, False) return lintersect_out + # Compute intersection point of two lines in 2D @field_operator def line_intersect( @@ -73,15 +108,18 @@ def line_intersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - - m1 = (line1_p2_lat - line1_p1_lat)/(line1_p2_lon - line1_p1_lon) - m2 = (line2_p2_lat - line2_p1_lat)/(line2_p2_lon - line2_p1_lon) - intersect_1 = (line2_p1_lat - line1_p1_lat + m1*line1_p1_lon - m2*line2_p1_lon) / (m1 - m2) - intersect_2 = line1_p1_lat + m1*(intersect_1 - line1_p1_lon) + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + intersect_1 = ( + line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon + ) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) return intersect_1, intersect_2 - + + @field_operator def _divide_flux_area_list_stencil_01( famask_int: Field[[EdgeDim, KDim], int32], @@ -97,18 +135,43 @@ def _divide_flux_area_list_stencil_01( dreg_patch0_3_y: Field[[EdgeDim, KDim], float], dreg_patch0_4_x: Field[[EdgeDim, KDim], float], dreg_patch0_4_y: Field[[EdgeDim, KDim], float], -) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: arrival_pts_1_x = dreg_patch0_1_x arrival_pts_1_y = dreg_patch0_1_y arrival_pts_2_x = dreg_patch0_2_x arrival_pts_2_y = dreg_patch0_2_y - depart_pts_1_x = dreg_patch0_4_x # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_x = dreg_patch0_4_x # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) depart_pts_1_y = dreg_patch0_4_y depart_pts_2_x = dreg_patch0_3_x depart_pts_2_y = dreg_patch0_3_y - lvn_pos = where(p_vn>=0.0, True, False) + lvn_pos = where(p_vn >= 0.0, True, False) # get flux area departure-line segment fl_line_p1_lon = depart_pts_1_x @@ -119,36 +182,98 @@ def _divide_flux_area_list_stencil_01( # get triangle edge 1 (A1V3) tri_line1_p1_lon = arrival_pts_1_x tri_line1_p1_lat = arrival_pts_1_y - tri_line1_p2_lon = where(lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim))) - tri_line1_p2_lat = where(lvn_pos, broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim))) + tri_line1_p2_lon = where( + lvn_pos, + broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim)), + ) + tri_line1_p2_lat = where( + lvn_pos, + broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim)), + ) # get triangle edge 2 (A2V3) tri_line2_p1_lon = arrival_pts_2_x tri_line2_p1_lat = arrival_pts_2_y - tri_line2_p2_lon = where(lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim))) - tri_line2_p2_lat = where(lvn_pos, broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim))) + tri_line2_p2_lon = where( + lvn_pos, + broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim)), + ) + tri_line2_p2_lat = where( + lvn_pos, + broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim)), + ) # Create first mask does departure-line segment intersects with A1V3 - lintersect_line1 = lintersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + lintersect_line1 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) # Create first mask does departure-line segment intersects with A2V3 - lintersect_line2 = lintersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + lintersect_line2 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) - lvn_sys_pos = where((p_vn * broadcast(tangent_orientation_dsl, (EdgeDim, KDim))) >= 0.0, True, False) + lvn_sys_pos = where( + (p_vn * broadcast(tangent_orientation_dsl, (EdgeDim, KDim))) >= 0.0, True, False + ) famask_bool = where(famask_int == int32(1), True, False) - #------------------------------------------------- Case 1 + # ------------------------------------------------- Case 1 mask_case1 = lintersect_line1 & lintersect_line2 & famask_bool - ps1_x, ps1_y = line_intersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) - ps2_x, ps2_y = line_intersect(fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, fl_line_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + ps1_x, ps1_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) # Case 1 - patch 0 dreg_patch0_1_x = where(mask_case1, arrival_pts_1_x, dreg_patch0_1_x) dreg_patch0_1_y = where(mask_case1, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where(mask_case1, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x) - dreg_patch0_2_y = where(mask_case1, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y) + dreg_patch0_2_x = where( + mask_case1, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x + ) + dreg_patch0_2_y = where( + mask_case1, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y + ) dreg_patch0_3_x = where(mask_case1, ps2_x, dreg_patch0_3_x) dreg_patch0_3_y = where(mask_case1, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where(mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x) - dreg_patch0_4_y = where(mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y) + dreg_patch0_4_x = where( + mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x + ) + dreg_patch0_4_y = where( + mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y + ) # Case 1 - patch 1 dreg_patch1_1_x = where(mask_case1, arrival_pts_1_x, 0.0) dreg_patch1_1_y = where(mask_case1, arrival_pts_1_y, 0.0) @@ -167,27 +292,43 @@ def _divide_flux_area_list_stencil_01( dreg_patch2_2_y = where(mask_case1, where(lvn_sys_pos, depart_pts_2_y, ps2_y), 0.0) dreg_patch2_3_x = where(mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_x), 0.0) dreg_patch2_3_y = where(mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_y), 0.0) - - #------------------------------------------------- Case 2a + + # ------------------------------------------------- Case 2a mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool # Case 2a - patch 0 dreg_patch0_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch0_1_x) dreg_patch0_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where(mask_case2a, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x) - dreg_patch0_2_y = where(mask_case2a, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y) + dreg_patch0_2_x = where( + mask_case2a, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x + ) + dreg_patch0_2_y = where( + mask_case2a, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y + ) dreg_patch0_3_x = where(mask_case2a, depart_pts_2_x, dreg_patch0_3_x) dreg_patch0_3_y = where(mask_case2a, depart_pts_2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where(mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x) - dreg_patch0_4_y = where(mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y) + dreg_patch0_4_x = where( + mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x + ) + dreg_patch0_4_y = where( + mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y + ) # Case 2a - patch 1 dreg_patch1_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_1_x) dreg_patch1_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_1_y) dreg_patch1_4_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_4_x) dreg_patch1_4_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_4_y) - dreg_patch1_2_x = where(mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_x), dreg_patch1_2_x) - dreg_patch1_2_y = where(mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_y), dreg_patch1_2_y) - dreg_patch1_3_x = where(mask_case2a, where(lvn_sys_pos, depart_pts_1_x, ps1_x), dreg_patch1_3_x) - dreg_patch1_3_y = where(mask_case2a, where(lvn_sys_pos, depart_pts_1_y, ps1_y), dreg_patch1_3_y) + dreg_patch1_2_x = where( + mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_x), dreg_patch1_2_x + ) + dreg_patch1_2_y = where( + mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_y), dreg_patch1_2_y + ) + dreg_patch1_3_x = where( + mask_case2a, where(lvn_sys_pos, depart_pts_1_x, ps1_x), dreg_patch1_3_x + ) + dreg_patch1_3_y = where( + mask_case2a, where(lvn_sys_pos, depart_pts_1_y, ps1_y), dreg_patch1_3_y + ) # Case 2a - patch 2 dreg_patch2_1_x = where(mask_case2a, 0.0, dreg_patch2_1_x) dreg_patch2_1_y = where(mask_case2a, 0.0, dreg_patch2_1_y) @@ -198,17 +339,33 @@ def _divide_flux_area_list_stencil_01( dreg_patch2_4_x = where(mask_case2a, 0.0, dreg_patch2_4_x) dreg_patch2_4_y = where(mask_case2a, 0.0, dreg_patch2_4_y) - #-------------------------------------------------- Case 2b + # -------------------------------------------------- Case 2b mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool # Case 2b - patch 0 dreg_patch0_1_x = where(mask_case2b, arrival_pts_1_x, dreg_patch0_1_x) dreg_patch0_1_y = where(mask_case2b, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where(mask_case2b, where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), dreg_patch0_2_x) - dreg_patch0_2_y = where(mask_case2b, where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), dreg_patch0_2_y) + dreg_patch0_2_x = where( + mask_case2b, + where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), + dreg_patch0_2_x, + ) + dreg_patch0_2_y = where( + mask_case2b, + where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), + dreg_patch0_2_y, + ) dreg_patch0_3_x = where(mask_case2b, ps2_x, dreg_patch0_3_x) dreg_patch0_3_y = where(mask_case2b, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where(mask_case2b, where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), dreg_patch0_4_x) - dreg_patch0_4_y = where(mask_case2b, where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), dreg_patch0_4_y) + dreg_patch0_4_x = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), + dreg_patch0_4_x, + ) + dreg_patch0_4_y = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), + dreg_patch0_4_y, + ) # Case 2b - patch 1 dreg_patch1_1_x = where(mask_case2b, 0.0, dreg_patch1_1_x) dreg_patch1_1_y = where(mask_case2b, 0.0, dreg_patch1_1_y) @@ -223,11 +380,19 @@ def _divide_flux_area_list_stencil_01( dreg_patch2_1_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_1_y) dreg_patch2_4_x = where(mask_case2b, arrival_pts_2_x, dreg_patch2_4_x) dreg_patch2_4_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_4_y) - dreg_patch2_2_x = where(mask_case2b, where(lvn_sys_pos, depart_pts_2_x, ps2_x), dreg_patch2_2_x) - dreg_patch2_2_y = where(mask_case2b, where(lvn_sys_pos, depart_pts_2_y, ps2_y), dreg_patch2_2_y) - dreg_patch2_3_x = where(mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_x), dreg_patch2_3_x) - dreg_patch2_3_y = where(mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_y), dreg_patch2_3_y) - + dreg_patch2_2_x = where( + mask_case2b, where(lvn_sys_pos, depart_pts_2_x, ps2_x), dreg_patch2_2_x + ) + dreg_patch2_2_y = where( + mask_case2b, where(lvn_sys_pos, depart_pts_2_y, ps2_y), dreg_patch2_2_y + ) + dreg_patch2_3_x = where( + mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_x), dreg_patch2_3_x + ) + dreg_patch2_3_y = where( + mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_y), dreg_patch2_3_y + ) + # flux area edge 1 and 2 fl_e1_p1_lon = arrival_pts_1_x fl_e1_p1_lat = arrival_pts_1_y @@ -238,29 +403,71 @@ def _divide_flux_area_list_stencil_01( fl_e2_p2_lon = depart_pts_2_x fl_e2_p2_lat = depart_pts_2_y - #----------------------------------------------- Case 3a + # ----------------------------------------------- Case 3a # Check whether flux area edge 2 intersects with triangle edge 1 - lintersect_e2_line1 = lintersect(fl_e2_p1_lon, fl_e2_p1_lat, fl_e2_p2_lon, fl_e2_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + lintersect_e2_line1 = lintersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) mask_case3a = lintersect_e2_line1 & famask_bool - pi1_x, pi1_y = line_intersect(fl_e2_p1_lon, fl_e2_p1_lat, fl_e2_p2_lon, fl_e2_p2_lat, tri_line1_p1_lon, tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat) + pi1_x, pi1_y = line_intersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) # Case 3a - patch 0 dreg_patch0_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch0_1_x) dreg_patch0_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where(mask_case3a, where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), dreg_patch0_2_x) - dreg_patch0_2_y = where(mask_case3a, where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), dreg_patch0_2_y) + dreg_patch0_2_x = where( + mask_case3a, + where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), + dreg_patch0_2_x, + ) + dreg_patch0_2_y = where( + mask_case3a, + where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), + dreg_patch0_2_y, + ) dreg_patch0_3_x = where(mask_case3a, ps2_x, dreg_patch0_3_x) dreg_patch0_3_y = where(mask_case3a, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), dreg_patch0_4_x) - dreg_patch0_4_y = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), dreg_patch0_4_y) + dreg_patch0_4_x = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), + dreg_patch0_4_x, + ) + dreg_patch0_4_y = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), + dreg_patch0_4_y, + ) # Case 3a - patch 1 dreg_patch1_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch1_1_x) dreg_patch1_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch1_1_y) - dreg_patch1_2_x = where(mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_x), dreg_patch1_2_x) - dreg_patch1_2_y = where(mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_y), dreg_patch1_2_y) + dreg_patch1_2_x = where( + mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_x), dreg_patch1_2_x + ) + dreg_patch1_2_y = where( + mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_y), dreg_patch1_2_y + ) dreg_patch1_3_x = where(mask_case3a, depart_pts_1_x, dreg_patch1_3_x) dreg_patch1_3_y = where(mask_case3a, depart_pts_1_y, dreg_patch1_3_y) - dreg_patch1_4_x = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_x, pi1_x), dreg_patch1_4_x) - dreg_patch1_4_y = where(mask_case3a, where(lvn_sys_pos, depart_pts_1_y, pi1_y), dreg_patch1_4_y) + dreg_patch1_4_x = where( + mask_case3a, where(lvn_sys_pos, depart_pts_1_x, pi1_x), dreg_patch1_4_x + ) + dreg_patch1_4_y = where( + mask_case3a, where(lvn_sys_pos, depart_pts_1_y, pi1_y), dreg_patch1_4_y + ) # Case 3a - patch 2 dreg_patch2_1_x = where(mask_case3a, 0.0, dreg_patch2_1_x) dreg_patch2_1_y = where(mask_case3a, 0.0, dreg_patch2_1_y) @@ -270,21 +477,47 @@ def _divide_flux_area_list_stencil_01( dreg_patch2_3_y = where(mask_case3a, 0.0, dreg_patch2_3_y) dreg_patch2_4_x = where(mask_case3a, 0.0, dreg_patch2_4_x) dreg_patch2_4_y = where(mask_case3a, 0.0, dreg_patch2_4_y) - - #------------------------------------------------ Case 3b + + # ------------------------------------------------ Case 3b # Check whether flux area edge 1 intersects with triangle edge 2 - lintersect_e1_line2 = lintersect(fl_e1_p1_lon, fl_e1_p1_lat, fl_e1_p2_lon, fl_e1_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + lintersect_e1_line2 = lintersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) mask_case3b = lintersect_e1_line2 & famask_bool - pi2_x, pi2_y = line_intersect(fl_e1_p1_lon, fl_e1_p1_lat, fl_e1_p2_lon, fl_e1_p2_lat, tri_line2_p1_lon, tri_line2_p1_lat, tri_line2_p2_lon, tri_line2_p2_lat) + pi2_x, pi2_y = line_intersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) # Case 3b - patch 0 dreg_patch0_1_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_1_x) dreg_patch0_1_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_1_y) dreg_patch0_4_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_4_x) dreg_patch0_4_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_4_y) - dreg_patch0_2_x = where(mask_case3b, where(lvn_sys_pos, arrival_pts_2_x, pi2_x), dreg_patch0_2_x) - dreg_patch0_2_y = where(mask_case3b, where(lvn_sys_pos, arrival_pts_2_y, pi2_y), dreg_patch0_2_y) - dreg_patch0_3_x = where(mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_x), dreg_patch0_3_x) - dreg_patch0_3_y = where(mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_y), dreg_patch0_3_y) + dreg_patch0_2_x = where( + mask_case3b, where(lvn_sys_pos, arrival_pts_2_x, pi2_x), dreg_patch0_2_x + ) + dreg_patch0_2_y = where( + mask_case3b, where(lvn_sys_pos, arrival_pts_2_y, pi2_y), dreg_patch0_2_y + ) + dreg_patch0_3_x = where( + mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_x), dreg_patch0_3_x + ) + dreg_patch0_3_y = where( + mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_y), dreg_patch0_3_y + ) # Case 3b - patch 1 dreg_patch1_1_x = where(mask_case3b, 0.0, dreg_patch1_1_x) dreg_patch1_1_y = where(mask_case3b, 0.0, dreg_patch1_1_y) @@ -297,18 +530,28 @@ def _divide_flux_area_list_stencil_01( # Case 3b - patch 2 dreg_patch2_1_x = where(mask_case3b, arrival_pts_2_x, dreg_patch2_1_x) dreg_patch2_1_y = where(mask_case3b, arrival_pts_2_y, dreg_patch2_1_y) - dreg_patch2_2_x = where(mask_case3b, where(lvn_sys_pos, depart_pts_2_x, pi2_x), dreg_patch2_2_x) - dreg_patch2_2_y = where(mask_case3b, where(lvn_sys_pos, depart_pts_2_y, pi2_y), dreg_patch2_2_y) + dreg_patch2_2_x = where( + mask_case3b, where(lvn_sys_pos, depart_pts_2_x, pi2_x), dreg_patch2_2_x + ) + dreg_patch2_2_y = where( + mask_case3b, where(lvn_sys_pos, depart_pts_2_y, pi2_y), dreg_patch2_2_y + ) dreg_patch2_3_x = where(mask_case3b, depart_pts_1_x, dreg_patch2_3_x) dreg_patch2_3_y = where(mask_case3b, depart_pts_1_y, dreg_patch2_3_y) - dreg_patch2_4_x = where(mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_x), dreg_patch2_4_x) - dreg_patch2_4_y = where(mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_y), dreg_patch2_4_y) + dreg_patch2_4_x = where( + mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_x), dreg_patch2_4_x + ) + dreg_patch2_4_y = where( + mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_y), dreg_patch2_4_y + ) - #--------------------------------------------- Case 4 + # --------------------------------------------- Case 4 # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions - # can't be overwritten by this new condition. - indices_previously_matched = (mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1) -# mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + # can't be overwritten by this new condition. + indices_previously_matched = ( + mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1 + ) + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices mask_case4 = famask_bool & (not indices_previously_matched) # Case 4 - patch 0 - no change # Case 4 - patch 1 @@ -330,8 +573,33 @@ def _divide_flux_area_list_stencil_01( dreg_patch2_4_x = where(mask_case4, 0.0, dreg_patch2_4_x) dreg_patch2_4_y = where(mask_case4, 0.0, dreg_patch2_4_y) + return ( + dreg_patch0_1_x, + dreg_patch0_1_y, + dreg_patch0_2_x, + dreg_patch0_2_y, + dreg_patch0_3_x, + dreg_patch0_3_y, + dreg_patch0_4_x, + dreg_patch0_4_y, + dreg_patch1_1_x, + dreg_patch1_1_y, + dreg_patch1_2_x, + dreg_patch1_2_y, + dreg_patch1_3_x, + dreg_patch1_3_y, + dreg_patch1_4_x, + dreg_patch1_4_y, + dreg_patch2_1_x, + dreg_patch2_1_y, + dreg_patch2_2_x, + dreg_patch2_2_y, + dreg_patch2_3_x, + dreg_patch2_3_y, + dreg_patch2_4_x, + dreg_patch2_4_y, + ) - return dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y @program def divide_flux_area_list_stencil_01( @@ -366,5 +634,43 @@ def divide_flux_area_list_stencil_01( dreg_patch2_4_y: Field[[EdgeDim, KDim], float], ): _divide_flux_area_list_stencil_01( - famask_int, p_vn, ptr_v3_lon, ptr_v3_lat, tangent_orientation_dsl, dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, out=(dreg_patch0_1_x, dreg_patch0_1_y, dreg_patch0_2_x, dreg_patch0_2_y, dreg_patch0_3_x, dreg_patch0_3_y, dreg_patch0_4_x, dreg_patch0_4_y, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y) + famask_int, + p_vn, + ptr_v3_lon, + ptr_v3_lat, + tangent_orientation_dsl, + dreg_patch0_1_x, + dreg_patch0_1_y, + dreg_patch0_2_x, + dreg_patch0_2_y, + dreg_patch0_3_x, + dreg_patch0_3_y, + dreg_patch0_4_x, + dreg_patch0_4_y, + out=( + dreg_patch0_1_x, + dreg_patch0_1_y, + dreg_patch0_2_x, + dreg_patch0_2_y, + dreg_patch0_3_x, + dreg_patch0_3_y, + dreg_patch0_4_x, + dreg_patch0_4_y, + dreg_patch1_1_x, + dreg_patch1_1_y, + dreg_patch1_2_x, + dreg_patch1_2_y, + dreg_patch1_3_x, + dreg_patch1_3_y, + dreg_patch1_4_x, + dreg_patch1_4_y, + dreg_patch2_1_x, + dreg_patch2_1_y, + dreg_patch2_2_x, + dreg_patch2_2_y, + dreg_patch2_3_x, + dreg_patch2_3_y, + dreg_patch2_4_x, + dreg_patch2_4_y, + ), ) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py index 84ce4cd376..8df569b509 100644 --- a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py @@ -11,11 +11,17 @@ # # 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, broadcast, int32, abs -from icon4py.common.dimension import E2EC, ECDim, CellDim, EdgeDim, KDim import sys + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim + + sys.setrecursionlimit(5500) + + @field_operator def _divide_flux_area_list_stencil_02( famask_int: Field[[EdgeDim, KDim], int32], @@ -48,14 +54,51 @@ def _divide_flux_area_list_stencil_02( dreg_patch2_3_y: Field[[EdgeDim, KDim], float], dreg_patch2_4_x: Field[[EdgeDim, KDim], float], dreg_patch2_4_y: Field[[EdgeDim, KDim], float], -) -> tuple[ Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32]]: +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], +]: famask_bool = where(famask_int == int32(1), True, False) lvn_pos = where(p_vn >= 0.0, True, False) # Translation of patch 1 and patch 2 in system relative to respective cell - bf_cc_patch1_lon = where(famask_bool, where(lvn_pos, bf_cc_patch1_lon(E2EC[0]), bf_cc_patch1_lon(E2EC[1])), 0.0) - bf_cc_patch1_lat = where(famask_bool, where(lvn_pos, bf_cc_patch1_lat(E2EC[0]), bf_cc_patch1_lat(E2EC[1])), 0.0) - bf_cc_patch2_lon = where(famask_bool, where(lvn_pos, bf_cc_patch2_lon(E2EC[0]), bf_cc_patch2_lon(E2EC[1])), 0.0) - bf_cc_patch2_lat = where(famask_bool, where(lvn_pos, bf_cc_patch2_lat(E2EC[0]), bf_cc_patch2_lat(E2EC[1])), 0.0) + bf_cc_patch1_lon = where( + famask_bool, + where(lvn_pos, bf_cc_patch1_lon(E2EC[0]), bf_cc_patch1_lon(E2EC[1])), + 0.0, + ) + bf_cc_patch1_lat = where( + famask_bool, + where(lvn_pos, bf_cc_patch1_lat(E2EC[0]), bf_cc_patch1_lat(E2EC[1])), + 0.0, + ) + bf_cc_patch2_lon = where( + famask_bool, + where(lvn_pos, bf_cc_patch2_lon(E2EC[0]), bf_cc_patch2_lon(E2EC[1])), + 0.0, + ) + bf_cc_patch2_lat = where( + famask_bool, + where(lvn_pos, bf_cc_patch2_lat(E2EC[0]), bf_cc_patch2_lat(E2EC[1])), + 0.0, + ) # patch1 in translated system dreg_patch1_1_x = dreg_patch1_1_x - bf_cc_patch1_lon @@ -78,20 +121,74 @@ def _divide_flux_area_list_stencil_02( # Store global index of the underlying grid cell # Adapt dimensions to fit ofr multiple levels - butterfly_idx_patch1_vnpos_3d = broadcast(butterfly_idx_patch1_vnpos, (EdgeDim, KDim)) - butterfly_idx_patch1_vnneg_3d = broadcast(butterfly_idx_patch1_vnneg, (EdgeDim, KDim)) - butterfly_idx_patch2_vnpos_3d = broadcast(butterfly_idx_patch2_vnpos, (EdgeDim, KDim)) - butterfly_idx_patch2_vnneg_3d = broadcast(butterfly_idx_patch2_vnneg, (EdgeDim, KDim)) - butterfly_blk_patch1_vnpos_3d = broadcast(butterfly_blk_patch1_vnpos, (EdgeDim, KDim)) - butterfly_blk_patch1_vnneg_3d = broadcast(butterfly_blk_patch1_vnneg, (EdgeDim, KDim)) - butterfly_blk_patch2_vnpos_3d = broadcast(butterfly_blk_patch2_vnpos, (EdgeDim, KDim)) - butterfly_blk_patch2_vnneg_3d = broadcast(butterfly_blk_patch2_vnneg, (EdgeDim, KDim)) - patch1_cell_idx_dsl = where(famask_bool, where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), int32(0)) - patch2_cell_idx_dsl = where(famask_bool, where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), int32(0)) - patch1_cell_blk_dsl = where(famask_bool, where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), int32(0)) - patch2_cell_blk_dsl = where(famask_bool, where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), int32(0)) + butterfly_idx_patch1_vnpos_3d = broadcast( + butterfly_idx_patch1_vnpos, (EdgeDim, KDim) + ) + butterfly_idx_patch1_vnneg_3d = broadcast( + butterfly_idx_patch1_vnneg, (EdgeDim, KDim) + ) + butterfly_idx_patch2_vnpos_3d = broadcast( + butterfly_idx_patch2_vnpos, (EdgeDim, KDim) + ) + butterfly_idx_patch2_vnneg_3d = broadcast( + butterfly_idx_patch2_vnneg, (EdgeDim, KDim) + ) + butterfly_blk_patch1_vnpos_3d = broadcast( + butterfly_blk_patch1_vnpos, (EdgeDim, KDim) + ) + butterfly_blk_patch1_vnneg_3d = broadcast( + butterfly_blk_patch1_vnneg, (EdgeDim, KDim) + ) + butterfly_blk_patch2_vnpos_3d = broadcast( + butterfly_blk_patch2_vnpos, (EdgeDim, KDim) + ) + butterfly_blk_patch2_vnneg_3d = broadcast( + butterfly_blk_patch2_vnneg, (EdgeDim, KDim) + ) + patch1_cell_idx_dsl = where( + famask_bool, + where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_dsl = where( + famask_bool, + where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_dsl = where( + famask_bool, + where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_dsl = where( + famask_bool, + where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) + + return ( + dreg_patch1_1_x, + dreg_patch1_1_y, + dreg_patch1_2_x, + dreg_patch1_2_y, + dreg_patch1_3_x, + dreg_patch1_3_y, + dreg_patch1_4_x, + dreg_patch1_4_y, + dreg_patch2_1_x, + dreg_patch2_1_y, + dreg_patch2_2_x, + dreg_patch2_2_y, + dreg_patch2_3_x, + dreg_patch2_3_y, + dreg_patch2_4_x, + dreg_patch2_4_y, + patch1_cell_idx_dsl, + patch1_cell_blk_dsl, + patch2_cell_idx_dsl, + patch2_cell_blk_dsl, + ) - return dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, patch1_cell_idx_dsl, patch1_cell_blk_dsl, patch2_cell_idx_dsl, patch2_cell_blk_dsl @program def divide_flux_area_list_stencil_02( @@ -131,5 +228,56 @@ def divide_flux_area_list_stencil_02( patch2_cell_blk_dsl: Field[[EdgeDim, KDim], int32], ): _divide_flux_area_list_stencil_02( - famask_int, p_vn, bf_cc_patch1_lon, bf_cc_patch1_lat, bf_cc_patch2_lon, bf_cc_patch2_lat, butterfly_idx_patch1_vnpos, butterfly_idx_patch1_vnneg, butterfly_blk_patch1_vnpos, butterfly_blk_patch1_vnneg, butterfly_idx_patch2_vnpos, butterfly_idx_patch2_vnneg, butterfly_blk_patch2_vnpos, butterfly_blk_patch2_vnneg, dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, out=( dreg_patch1_1_x, dreg_patch1_1_y, dreg_patch1_2_x, dreg_patch1_2_y, dreg_patch1_3_x, dreg_patch1_3_y, dreg_patch1_4_x, dreg_patch1_4_y, dreg_patch2_1_x, dreg_patch2_1_y, dreg_patch2_2_x, dreg_patch2_2_y, dreg_patch2_3_x, dreg_patch2_3_y, dreg_patch2_4_x, dreg_patch2_4_y, patch1_cell_idx_dsl, patch1_cell_blk_dsl, patch2_cell_idx_dsl, patch2_cell_blk_dsl) + famask_int, + p_vn, + bf_cc_patch1_lon, + bf_cc_patch1_lat, + bf_cc_patch2_lon, + bf_cc_patch2_lat, + butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg, + dreg_patch1_1_x, + dreg_patch1_1_y, + dreg_patch1_2_x, + dreg_patch1_2_y, + dreg_patch1_3_x, + dreg_patch1_3_y, + dreg_patch1_4_x, + dreg_patch1_4_y, + dreg_patch2_1_x, + dreg_patch2_1_y, + dreg_patch2_2_x, + dreg_patch2_2_y, + dreg_patch2_3_x, + dreg_patch2_3_y, + dreg_patch2_4_x, + dreg_patch2_4_y, + out=( + dreg_patch1_1_x, + dreg_patch1_1_y, + dreg_patch1_2_x, + dreg_patch1_2_y, + dreg_patch1_3_x, + dreg_patch1_3_y, + dreg_patch1_4_x, + dreg_patch1_4_y, + dreg_patch2_1_x, + dreg_patch2_1_y, + dreg_patch2_2_x, + dreg_patch2_2_y, + dreg_patch2_3_x, + dreg_patch2_3_y, + dreg_patch2_4_x, + dreg_patch2_4_y, + patch1_cell_idx_dsl, + patch1_cell_blk_dsl, + patch2_cell_idx_dsl, + patch2_cell_blk_dsl, + ), ) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py index c78e704b1e..53945a844e 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs, int32, where +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim, Koff +from icon4py.common.dimension import CellDim, KDim, Koff @field_operator @@ -25,7 +25,13 @@ def _face_val_ppm_stencil_01a( zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) zfac = (p_cc(Koff[+1]) - p_cc) / (p_cellhgt_mc_now(Koff[+1]) + p_cellhgt_mc_now) - z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[+1])) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now(Koff[+1])) * zfac_m1) + z_slope = ( + p_cellhgt_mc_now + / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[+1])) + ) * ( + (2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + + (p_cellhgt_mc_now + 2.0 * p_cellhgt_mc_now(Koff[+1])) * zfac_m1 + ) return z_slope @@ -38,7 +44,13 @@ def _face_val_ppm_stencil_01b( zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) - z_slope = ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) ) * ( (2.*p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.*p_cellhgt_mc_now) * zfac_m1) + z_slope = ( + p_cellhgt_mc_now + / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) + ) * ( + (2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + + (p_cellhgt_mc_now + 2.0 * p_cellhgt_mc_now) * zfac_m1 + ) return z_slope @@ -53,9 +65,11 @@ def _face_val_ppm_stencil_01( vert_idx = broadcast(vert_idx, (CellDim, KDim)) - z_slope = where(vert_idx Field[[CellDim, KDim], float]: - p_face = p_cc*(1. - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + (p_cellhgt_mc_now/(p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))* p_cc + p_cc(Koff[-1])) + p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) return p_face @@ -36,6 +38,7 @@ def _face_val_ppm_stencil_02b( p_face = p_cc return p_face + @field_operator def _face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], @@ -54,16 +57,20 @@ def _face_val_ppm_stencil_02( slev: int32, elev: int32, slevp1: int32, - elevp1: int32 + elevp1: int32, ) -> Field[[CellDim, KDim], float]: vert_idx = broadcast(vert_idx, (CellDim, KDim)) - p_face = where( (vert_idx==slevp1) | (vert_idx==elev), _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), p_face_in) + p_face = where( + (vert_idx == slevp1) | (vert_idx == elev), + _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), + p_face_in, + ) - p_face = where( (vert_idx==slev), _face_val_ppm_stencil_02b(p_cc), p_face) + p_face = where((vert_idx == slev), _face_val_ppm_stencil_02b(p_cc), p_face) - p_face = where( (vert_idx==elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) + p_face = where((vert_idx == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) return p_face @@ -90,4 +97,4 @@ def face_val_ppm_stencil_02( slevp1, elevp1, out=p_face, - ) \ No newline at end of file + ) diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py index 5576d3ad83..70684d4baf 100644 --- a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py +++ b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py @@ -13,6 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field + from icon4py.common.dimension import EdgeDim, KDim @@ -20,7 +21,7 @@ def _hflux_ffsl_hybrid_stencil_02( p_out_e: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], - z_dreg_area: Field[[EdgeDim, KDim], float] + z_dreg_area: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: p_out_e = p_mass_flx_e * p_out_e / z_dreg_area @@ -32,7 +33,7 @@ def _hflux_ffsl_hybrid_stencil_02( def hflux_ffsl_hybrid_stencil_02( p_out_e: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], - z_dreg_area: Field[[EdgeDim, KDim], float] + z_dreg_area: Field[[EdgeDim, KDim], float], ): _hflux_ffsl_hybrid_stencil_02( p_out_e, diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py index 727dcd83bc..cf86fbce4f 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, abs +from gt4py.next.ffront.fbuiltins import Field, abs -from icon4py.common.dimension import C2CE, C2E, CEDim, E2C, CellDim, EdgeDim, KDim +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -26,8 +26,11 @@ def _hflx_limiter_mo_stencil_01a( p_cc: Field[[CellDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) - abs(p_mass_flx_e)* ( p_cc(E2C[0]) - p_cc(E2C[1]) ) ) - # z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) ) + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc(E2C[0]) + p_cc(E2C[1])) + - abs(p_mass_flx_e) * (p_cc(E2C[0]) - p_cc(E2C[1])) + ) + # this is comented out: z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) ) z_anti = p_mflx_tracer_h - z_mflx_low return (z_anti, z_mflx_low) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py index 689e6dd8ad..ada965d084 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py @@ -12,9 +12,15 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, neighbor_sum +from gt4py.next.ffront.fbuiltins import ( + Field, + broadcast, + maximum, + minimum, + neighbor_sum, +) -from icon4py.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim, C2EDim, E2C +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -27,19 +33,31 @@ def _hflx_limiter_mo_stencil_01b( p_cc: Field[[CellDim, KDim], float], z_fluxdiv_c: Field[[CellDim, KDim], float], p_dtime: float, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: z_mflx_anti = broadcast(p_rhodz_new, (CellDim, C2EDim, KDim)) - z_mflx_anti_in = -1. * neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) - z_mflx_anti_out = neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) - z_fluxdiv_c = neighbor_sum( z_mflx_low(C2E) * geofac_div, axis=C2EDim) + z_mflx_anti_in = -1.0 * neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) + z_mflx_anti_out = neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) + z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div, axis=C2EDim) - z_tracer_new_low = ( p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c ) / p_rhodz_new - z_tracer_max = maximum(p_cc,z_tracer_new_low) - z_tracer_min = minimum(p_cc,z_tracer_new_low) + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = maximum(p_cc, z_tracer_new_low) + z_tracer_min = minimum(p_cc, z_tracer_new_low) - return (z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, z_tracer_max, z_tracer_min) + return ( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ) @program @@ -67,5 +85,11 @@ def hflx_limiter_mo_stencil_01b( p_cc, z_fluxdiv_c, p_dtime, - out=(z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, z_tracer_max, z_tracer_min), + out=( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ), ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py index 2160275ee7..cf277deafe 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py @@ -27,7 +27,11 @@ def _hflx_limiter_mo_stencil_02( z_tracer_min: Field[[CellDim, KDim], float], lo_bound: int32, hi_bound: int32, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound) z_tracer_new_out = where( condition, @@ -38,7 +42,7 @@ def _hflx_limiter_mo_stencil_02( z_tracer_max_out = where(condition, maximum(p_cc, z_tracer_new_out), z_tracer_max) z_tracer_min_out = where(condition, minimum(p_cc, z_tracer_new_out), z_tracer_min) - return(z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) + return (z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) @program diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py index a41ae7913c..3abad684d3 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py @@ -56,7 +56,6 @@ def _hflx_limiter_mo_stencil_03a( z_max: Field[[CellDim, KDim], float], z_min: Field[[CellDim, KDim], float], dbl_eps: float, - ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) @@ -75,13 +74,13 @@ def _hflx_limiter_mo_stencil_03( z_mflx_anti_out: Field[[CellDim, KDim], float], z_tracer_new_low: Field[[CellDim, KDim], float], dbl_eps: float, - ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_max, z_min = _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) - z_max, z_min = _hflx_limiter_mo_stencil_03_min_max( z_tracer_max, z_tracer_min, beta_fct, r_beta_fct ) - - r_p, r_m =_hflx_limiter_mo_stencil_03a( + r_p, r_m = _hflx_limiter_mo_stencil_03a( z_mflx_anti_in, z_mflx_anti_out, z_tracer_new_low, @@ -91,6 +90,7 @@ def _hflx_limiter_mo_stencil_03( ) return r_p, r_m + @program def hflx_limiter_mo_stencil_03( z_tracer_max: Field[[CellDim, KDim], float], @@ -115,4 +115,4 @@ def hflx_limiter_mo_stencil_03( z_tracer_new_low, dbl_eps, out=(r_p, r_m), - ) \ No newline at end of file + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py index d5f342766d..2e8e82c4f9 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py @@ -13,9 +13,9 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import max_over, maximum, min_over, minimum, where +from gt4py.next.ffront.fbuiltins import minimum, where -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, EdgeDim, E2C +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -26,16 +26,14 @@ def _hflx_limiter_mo_stencil_05( r_p: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_signum = where( (z_anti>0.), 1., -1.) + z_signum = where((z_anti > 0.0), 1.0, -1.0) - r_frac = 0.5*( (1. + z_signum)* - minimum(r_m(E2C[0]), - r_p(E2C[1]) ) - + (1.-z_signum)* - minimum( r_m(E2C[1]), - r_p(E2C[0]) ) ) + r_frac = 0.5 * ( + (1.0 + z_signum) * minimum(r_m(E2C[0]), r_p(E2C[1])) + + (1.0 - z_signum) * minimum(r_m(E2C[1]), r_p(E2C[0])) + ) - p_mflx_tracer_h = z_mflx_low + minimum(1.,r_frac) * z_anti + p_mflx_tracer_h = z_mflx_low + minimum(1.0, r_frac) * z_anti return p_mflx_tracer_h @@ -48,11 +46,4 @@ def hflx_limiter_mo_stencil_05( r_p: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], ): - _hflx_limiter_mo_stencil_05( - z_anti, - z_mflx_low, - r_m, - r_p, - out=p_mflx_tracer_h - ) - + _hflx_limiter_mo_stencil_05(z_anti, z_mflx_low, r_m, r_p, out=p_mflx_tracer_h) diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py index 5de0f465a8..37119dcc76 100644 --- a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.ffront.fbuiltins import Field, maximum, abs, where, broadcast, int32 from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.common.dimension import EdgeDim, KDim @@ -59,91 +59,271 @@ def _prep_gauss_quadrature_c_list_stencil( dbl_eps: float, eps: float, p_dreg_area: Field[[EdgeDim, KDim], float], -) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float] +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], ]: z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - + z_eta_1_1 = 1.0 - eta_1 - z_eta_2_1 = 1.0 - eta_2 - z_eta_3_1 = 1.0 - eta_3 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 z_eta_4_1 = 1.0 - eta_4 z_eta_1_2 = 1.0 + eta_1 - z_eta_2_2 = 1.0 + eta_2 - z_eta_3_2 = 1.0 + eta_3 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 z_eta_4_2 = 1.0 + eta_4 z_eta_1_3 = 1.0 - zeta_1 - z_eta_2_3 = 1.0 - zeta_2 - z_eta_3_3 = 1.0 - zeta_3 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 z_eta_4_3 = 1.0 - zeta_4 z_eta_1_4 = 1.0 + zeta_1 - z_eta_2_4 = 1.0 + zeta_2 - z_eta_3_4 = 1.0 + zeta_3 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 z_eta_4_4 = 1.0 + zeta_4 famask_bool = where(famask_int == int32(1), True, False) - p_coords_dreg_v_1_x = where(famask_bool, p_coords_dreg_v_1_x, 0.0) - p_coords_dreg_v_2_x = where(famask_bool, p_coords_dreg_v_2_x, 0.0) - p_coords_dreg_v_3_x = where(famask_bool, p_coords_dreg_v_3_x, 0.0) - p_coords_dreg_v_4_x = where(famask_bool, p_coords_dreg_v_4_x, 0.0) - p_coords_dreg_v_1_y = where(famask_bool, p_coords_dreg_v_1_y, 0.0) - p_coords_dreg_v_2_y = where(famask_bool, p_coords_dreg_v_2_y, 0.0) - p_coords_dreg_v_3_y = where(famask_bool, p_coords_dreg_v_3_y, 0.0) - p_coords_dreg_v_4_y = where(famask_bool, p_coords_dreg_v_4_y, 0.0) + p_coords_dreg_v_1_x = where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = where(famask_bool, p_coords_dreg_v_4_y, 0.0) - wgt_t_detjac_1 = where(famask_bool, dbl_eps + z_wgt_1 * ( - ( z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) - wgt_t_detjac_2 = where(famask_bool, dbl_eps + z_wgt_2 * ( - ( z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) - wgt_t_detjac_3 = where(famask_bool, dbl_eps + z_wgt_3 * ( - ( z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) - wgt_t_detjac_4 = where(famask_bool, dbl_eps + z_wgt_4 * ( - ( z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))), 0.0) + wgt_t_detjac_1 = where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_2 = where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_3 = where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 = where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) - z_gauss_pts_1_x = shape_func_1_1 * p_coords_dreg_v_1_x + shape_func_2_1 * p_coords_dreg_v_2_x + shape_func_3_1 * p_coords_dreg_v_3_x + shape_func_4_1 * p_coords_dreg_v_4_x - z_gauss_pts_1_y = shape_func_1_1 * p_coords_dreg_v_1_y + shape_func_2_1 * p_coords_dreg_v_2_y + shape_func_3_1 * p_coords_dreg_v_3_y + shape_func_4_1 * p_coords_dreg_v_4_y - z_gauss_pts_2_x = shape_func_1_2 * p_coords_dreg_v_1_x + shape_func_2_2 * p_coords_dreg_v_2_x + shape_func_3_2 * p_coords_dreg_v_3_x + shape_func_4_2 * p_coords_dreg_v_4_x - z_gauss_pts_2_y = shape_func_1_2 * p_coords_dreg_v_1_y + shape_func_2_2 * p_coords_dreg_v_2_y + shape_func_3_2 * p_coords_dreg_v_3_y + shape_func_4_2 * p_coords_dreg_v_4_y - z_gauss_pts_3_x = shape_func_1_3 * p_coords_dreg_v_1_x + shape_func_2_3 * p_coords_dreg_v_2_x + shape_func_3_3 * p_coords_dreg_v_3_x + shape_func_4_3 * p_coords_dreg_v_4_x - z_gauss_pts_3_y = shape_func_1_3 * p_coords_dreg_v_1_y + shape_func_2_3 * p_coords_dreg_v_2_y + shape_func_3_3 * p_coords_dreg_v_3_y + shape_func_4_3 * p_coords_dreg_v_4_y - z_gauss_pts_4_x = shape_func_1_4 * p_coords_dreg_v_1_x + shape_func_2_4 * p_coords_dreg_v_2_x + shape_func_3_4 * p_coords_dreg_v_3_x + shape_func_4_4 * p_coords_dreg_v_4_x - z_gauss_pts_4_y = shape_func_1_4 * p_coords_dreg_v_1_y + shape_func_2_4 * p_coords_dreg_v_2_y + shape_func_3_4 * p_coords_dreg_v_3_y + shape_func_4_4 * p_coords_dreg_v_4_y + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x + wgt_t_detjac_3 * z_gauss_pts_3_x + wgt_t_detjac_4 * z_gauss_pts_4_x - p_quad_vector_sum_3 = wgt_t_detjac_1 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_y - p_quad_vector_sum_4 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + wgt_t_detjac_4 * z_gauss_pts_4_x**2 - p_quad_vector_sum_5 = wgt_t_detjac_1 * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_y**2 - p_quad_vector_sum_6 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - p_quad_vector_sum_7 = wgt_t_detjac_1 * z_gauss_pts_1_x**3 + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + wgt_t_detjac_4 * z_gauss_pts_4_x**3 - p_quad_vector_sum_8 = wgt_t_detjac_1 * z_gauss_pts_1_y**3 + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + wgt_t_detjac_4 * z_gauss_pts_4_y**3 - p_quad_vector_sum_9 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y - p_quad_vector_sum_10 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + p_quad_vector_sum_1 = ( + wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + ) + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**3 + + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**3 + + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + ) p_dreg_area = p_dreg_area + p_quad_vector_sum_1 - - return p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3,p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10 + return ( + p_dreg_area, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) @program @@ -200,5 +380,57 @@ def prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], ): _prep_gauss_quadrature_c_list_stencil( - famask_int, p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, p_coords_dreg_v_4_x, p_coords_dreg_v_1_y, p_coords_dreg_v_2_y, p_coords_dreg_v_3_y, p_coords_dreg_v_4_y, shape_func_1_1, shape_func_2_1, shape_func_3_1, shape_func_4_1, shape_func_1_2, shape_func_2_2, shape_func_3_2, shape_func_4_2, shape_func_1_3, shape_func_2_3, shape_func_3_3, shape_func_4_3, shape_func_1_4, shape_func_2_4, shape_func_3_4, shape_func_4_4, zeta_1, zeta_2, zeta_3, zeta_4, eta_1, eta_2, eta_3, eta_4, wgt_zeta_1, wgt_zeta_2, wgt_eta_1, wgt_eta_2, dbl_eps, eps, p_dreg_area, out=(p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10) + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_dreg_area, + out=( + p_dreg_area, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ), ) diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py b/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py index 8b68ada345..6aa36aab23 100644 --- a/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py +++ b/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.ffront.fbuiltins import Field, maximum, abs, where, broadcast from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, abs, maximum, where from icon4py.common.dimension import EdgeDim, KDim @@ -57,82 +57,240 @@ def _prep_gauss_quadrature_c_stencil( wgt_eta_2: float, dbl_eps: float, eps: float, -) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float] +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], ]: z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - + z_eta_1_1 = 1.0 - eta_1 - z_eta_2_1 = 1.0 - eta_2 - z_eta_3_1 = 1.0 - eta_3 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 z_eta_4_1 = 1.0 - eta_4 z_eta_1_2 = 1.0 + eta_1 - z_eta_2_2 = 1.0 + eta_2 - z_eta_3_2 = 1.0 + eta_3 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 z_eta_4_2 = 1.0 + eta_4 z_eta_1_3 = 1.0 - zeta_1 - z_eta_2_3 = 1.0 - zeta_2 - z_eta_3_3 = 1.0 - zeta_3 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 z_eta_4_3 = 1.0 - zeta_4 z_eta_1_4 = 1.0 + zeta_1 - z_eta_2_4 = 1.0 + zeta_2 - z_eta_3_4 = 1.0 + zeta_3 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 z_eta_4_4 = 1.0 + zeta_4 wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( - ( z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( - ( z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( - ( z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( - ( z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x)) - * (z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y)) - - (z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y)) - * (z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x))) + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) - z_gauss_pts_1_x = shape_func_1_1 * p_coords_dreg_v_1_x + shape_func_2_1 * p_coords_dreg_v_2_x + shape_func_3_1 * p_coords_dreg_v_3_x + shape_func_4_1 * p_coords_dreg_v_4_x - z_gauss_pts_1_y = shape_func_1_1 * p_coords_dreg_v_1_y + shape_func_2_1 * p_coords_dreg_v_2_y + shape_func_3_1 * p_coords_dreg_v_3_y + shape_func_4_1 * p_coords_dreg_v_4_y - z_gauss_pts_2_x = shape_func_1_2 * p_coords_dreg_v_1_x + shape_func_2_2 * p_coords_dreg_v_2_x + shape_func_3_2 * p_coords_dreg_v_3_x + shape_func_4_2 * p_coords_dreg_v_4_x - z_gauss_pts_2_y = shape_func_1_2 * p_coords_dreg_v_1_y + shape_func_2_2 * p_coords_dreg_v_2_y + shape_func_3_2 * p_coords_dreg_v_3_y + shape_func_4_2 * p_coords_dreg_v_4_y - z_gauss_pts_3_x = shape_func_1_3 * p_coords_dreg_v_1_x + shape_func_2_3 * p_coords_dreg_v_2_x + shape_func_3_3 * p_coords_dreg_v_3_x + shape_func_4_3 * p_coords_dreg_v_4_x - z_gauss_pts_3_y = shape_func_1_3 * p_coords_dreg_v_1_y + shape_func_2_3 * p_coords_dreg_v_2_y + shape_func_3_3 * p_coords_dreg_v_3_y + shape_func_4_3 * p_coords_dreg_v_4_y - z_gauss_pts_4_x = shape_func_1_4 * p_coords_dreg_v_1_x + shape_func_2_4 * p_coords_dreg_v_2_x + shape_func_3_4 * p_coords_dreg_v_3_x + shape_func_4_4 * p_coords_dreg_v_4_x - z_gauss_pts_4_y = shape_func_1_4 * p_coords_dreg_v_1_y + shape_func_2_4 * p_coords_dreg_v_2_y + shape_func_3_4 * p_coords_dreg_v_3_y + shape_func_4_4 * p_coords_dreg_v_4_y + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x + wgt_t_detjac_3 * z_gauss_pts_3_x + wgt_t_detjac_4 * z_gauss_pts_4_x - p_quad_vector_sum_3 = wgt_t_detjac_1 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_y - p_quad_vector_sum_4 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + wgt_t_detjac_4 * z_gauss_pts_4_x**2 - p_quad_vector_sum_5 = wgt_t_detjac_1 * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_y**2 - p_quad_vector_sum_6 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - p_quad_vector_sum_7 = wgt_t_detjac_1 * z_gauss_pts_1_x**3 + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + wgt_t_detjac_4 * z_gauss_pts_4_x**3 - p_quad_vector_sum_8 = wgt_t_detjac_1 * z_gauss_pts_1_y**3 + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + wgt_t_detjac_4 * z_gauss_pts_4_y**3 - p_quad_vector_sum_9 = wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y - p_quad_vector_sum_10 = wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + p_quad_vector_sum_1 = ( + wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + ) + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**3 + + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**3 + + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + ) z_area = p_quad_vector_sum_1 - p_dreg_area_out = where(z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area))) - + p_dreg_area_out = where( + z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area)) + ) - return p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3,p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6,p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9,p_quad_vector_sum_10, p_dreg_area_out + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ) @program @@ -188,5 +346,55 @@ def prep_gauss_quadrature_c_stencil( p_dreg_area_out: Field[[EdgeDim, KDim], float], ): _prep_gauss_quadrature_c_stencil( - p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, p_coords_dreg_v_4_x, p_coords_dreg_v_1_y, p_coords_dreg_v_2_y, p_coords_dreg_v_3_y, p_coords_dreg_v_4_y, shape_func_1_1, shape_func_2_1, shape_func_3_1, shape_func_4_1, shape_func_1_2, shape_func_2_2, shape_func_3_2, shape_func_4_2, shape_func_1_3, shape_func_2_3, shape_func_3_3, shape_func_4_3, shape_func_1_4, shape_func_2_4, shape_func_3_4, shape_func_4_4, zeta_1, zeta_2, zeta_3, zeta_4, eta_1, eta_2, eta_3, eta_4, wgt_zeta_1, wgt_zeta_2, wgt_eta_1, wgt_eta_2, dbl_eps, eps, out=(p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, p_quad_vector_sum_4, p_quad_vector_sum_5, p_quad_vector_sum_6, p_quad_vector_sum_7, p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, p_dreg_area_out) + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + out=( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ), ) diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py index dedda2ce10..fdfb97fb8e 100644 --- a/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py +++ b/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py @@ -11,14 +11,18 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import sys # Increase recusion depth, otherwise it doesn't compile + from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast + from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim -import sys # Increase recusion depth, otherwise it doesn't compile + sys.setrecursionlimit(6000) + @field_operator def _recon_lsq_cell_c_stencil( p_cc: Field[[CellDim, KDim], float], @@ -31,15 +35,15 @@ def _recon_lsq_cell_c_stencil( lsq_qtmat_c_7: Field[[CECECDim], float], lsq_qtmat_c_8: Field[[CECECDim], float], lsq_qtmat_c_9: Field[[CECECDim], float], - lsq_rmat_rdiag_c_1: Field[[CellDim], float], + lsq_rmat_rdiag_c_1: Field[[CellDim], float], lsq_rmat_rdiag_c_2: Field[[CellDim], float], - lsq_rmat_rdiag_c_3: Field[[CellDim], float], - lsq_rmat_rdiag_c_4: Field[[CellDim], float], - lsq_rmat_rdiag_c_5: Field[[CellDim], float], - lsq_rmat_rdiag_c_6: Field[[CellDim], float], - lsq_rmat_rdiag_c_7: Field[[CellDim], float], - lsq_rmat_rdiag_c_8: Field[[CellDim], float], - lsq_rmat_rdiag_c_9: Field[[CellDim], float], + lsq_rmat_rdiag_c_3: Field[[CellDim], float], + lsq_rmat_rdiag_c_4: Field[[CellDim], float], + lsq_rmat_rdiag_c_5: Field[[CellDim], float], + lsq_rmat_rdiag_c_6: Field[[CellDim], float], + lsq_rmat_rdiag_c_7: Field[[CellDim], float], + lsq_rmat_rdiag_c_8: Field[[CellDim], float], + lsq_rmat_rdiag_c_9: Field[[CellDim], float], lsq_rmat_utri_c_1: Field[[CellDim], float], lsq_rmat_utri_c_2: Field[[CellDim], float], lsq_rmat_utri_c_3: Field[[CellDim], float], @@ -85,7 +89,18 @@ def _recon_lsq_cell_c_stencil( lsq_moments_7: Field[[CellDim], float], lsq_moments_8: Field[[CellDim], float], lsq_moments_9: Field[[CellDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: lsq_rmat_rdiag_c_1 = broadcast(lsq_rmat_rdiag_c_1, (CellDim, KDim)) lsq_rmat_rdiag_c_2 = broadcast(lsq_rmat_rdiag_c_2, (CellDim, KDim)) lsq_rmat_rdiag_c_3 = broadcast(lsq_rmat_rdiag_c_3, (CellDim, KDim)) @@ -105,26 +120,185 @@ def _recon_lsq_cell_c_stencil( lsq_qtmat_c_8 = broadcast(lsq_qtmat_c_8, (CECECDim, KDim)) lsq_qtmat_c_9 = broadcast(lsq_qtmat_c_9, (CECECDim, KDim)) - p_coeff_10 = lsq_rmat_rdiag_c_9 * (lsq_qtmat_c_9(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_9(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_9(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_9(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_9(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_9(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_9(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_9(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_9(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc)) + p_coeff_10 = lsq_rmat_rdiag_c_9 * ( + lsq_qtmat_c_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_9(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + + p_coeff_9 = lsq_rmat_rdiag_c_8 * ( + lsq_qtmat_c_8(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_8(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - lsq_rmat_utri_c_1 * p_coeff_10 + ) - p_coeff_9 = lsq_rmat_rdiag_c_8 * (lsq_qtmat_c_8(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_8(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_8(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_8(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_8(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_8(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_8(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_8(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_8(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - lsq_rmat_utri_c_1*p_coeff_10) + p_coeff_8 = lsq_rmat_rdiag_c_7 * ( + lsq_qtmat_c_7(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_7(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) + ) - p_coeff_8 = lsq_rmat_rdiag_c_7 * (lsq_qtmat_c_7(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_7(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_7(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_7(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_7(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_7(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_7(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_7(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_7(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_2*p_coeff_9 + lsq_rmat_utri_c_3*p_coeff_10)) + p_coeff_7 = lsq_rmat_rdiag_c_6 * ( + lsq_qtmat_c_6(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_6(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_4 * p_coeff_8 + + lsq_rmat_utri_c_5 * p_coeff_9 + + lsq_rmat_utri_c_6 * p_coeff_10 + ) + ) - p_coeff_7 = lsq_rmat_rdiag_c_6 * (lsq_qtmat_c_6(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_6(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_6(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_6(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_6(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_6(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_6(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_6(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_6(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_4*p_coeff_8 + lsq_rmat_utri_c_5*p_coeff_9 + lsq_rmat_utri_c_6*p_coeff_10)) + p_coeff_6 = lsq_rmat_rdiag_c_5 * ( + lsq_qtmat_c_5(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_5(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_7 * p_coeff_7 + + lsq_rmat_utri_c_8 * p_coeff_8 + + lsq_rmat_utri_c_9 * p_coeff_9 + + lsq_rmat_utri_c_10 * p_coeff_10 + ) + ) - p_coeff_6 = lsq_rmat_rdiag_c_5 * (lsq_qtmat_c_5(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_5(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_5(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_5(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_5(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_5(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_5(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_5(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_5(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_7*p_coeff_7 + lsq_rmat_utri_c_8*p_coeff_8 + lsq_rmat_utri_c_9*p_coeff_9 + lsq_rmat_utri_c_10*p_coeff_10)) + p_coeff_5 = lsq_rmat_rdiag_c_4 * ( + lsq_qtmat_c_4(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_4(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_11 * p_coeff_6 + + lsq_rmat_utri_c_12 * p_coeff_7 + + lsq_rmat_utri_c_13 * p_coeff_8 + + lsq_rmat_utri_c_14 * p_coeff_9 + + lsq_rmat_utri_c_15 * p_coeff_10 + ) + ) - p_coeff_5 = lsq_rmat_rdiag_c_4 * (lsq_qtmat_c_4(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_4(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_4(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_4(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_4(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_4(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_4(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_4(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_4(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_11*p_coeff_6 + lsq_rmat_utri_c_12*p_coeff_7 + lsq_rmat_utri_c_13*p_coeff_8 + lsq_rmat_utri_c_14*p_coeff_9 + lsq_rmat_utri_c_15*p_coeff_10)) + p_coeff_4 = lsq_rmat_rdiag_c_3 * ( + lsq_qtmat_c_3(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_3(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_16 * p_coeff_5 + + lsq_rmat_utri_c_17 * p_coeff_6 + + lsq_rmat_utri_c_18 * p_coeff_7 + + lsq_rmat_utri_c_19 * p_coeff_8 + + lsq_rmat_utri_c_20 * p_coeff_9 + + lsq_rmat_utri_c_21 * p_coeff_10 + ) + ) - p_coeff_4 = lsq_rmat_rdiag_c_3 * (lsq_qtmat_c_3(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_3(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_3(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_3(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_3(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_3(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_3(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_3(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_3(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_16*p_coeff_5 + lsq_rmat_utri_c_17*p_coeff_6 + lsq_rmat_utri_c_18*p_coeff_7 + lsq_rmat_utri_c_19*p_coeff_8 + lsq_rmat_utri_c_20*p_coeff_9 + lsq_rmat_utri_c_21*p_coeff_10)) + p_coeff_3 = lsq_rmat_rdiag_c_2 * ( + lsq_qtmat_c_2(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_2(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_22 * p_coeff_4 + + lsq_rmat_utri_c_23 * p_coeff_5 + + lsq_rmat_utri_c_24 * p_coeff_6 + + lsq_rmat_utri_c_25 * p_coeff_7 + + lsq_rmat_utri_c_26 * p_coeff_8 + + lsq_rmat_utri_c_27 * p_coeff_9 + + lsq_rmat_utri_c_28 * p_coeff_10 + ) + ) - p_coeff_3 = lsq_rmat_rdiag_c_2 * (lsq_qtmat_c_2(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_2(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_2(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_2(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_2(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_2(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_2(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_2(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_2(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_22*p_coeff_4 + lsq_rmat_utri_c_23*p_coeff_5 + lsq_rmat_utri_c_24*p_coeff_6 + lsq_rmat_utri_c_25*p_coeff_7 + lsq_rmat_utri_c_26*p_coeff_8 + lsq_rmat_utri_c_27*p_coeff_9 + lsq_rmat_utri_c_28*p_coeff_10)) + p_coeff_2 = lsq_rmat_rdiag_c_1 * ( + lsq_qtmat_c_1(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_qtmat_c_1(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + - ( + lsq_rmat_utri_c_29 * p_coeff_3 + + lsq_rmat_utri_c_30 * p_coeff_4 + + lsq_rmat_utri_c_31 * p_coeff_5 + + lsq_rmat_utri_c_32 * p_coeff_6 + + lsq_rmat_utri_c_33 * p_coeff_7 + + lsq_rmat_utri_c_34 * p_coeff_8 + + lsq_rmat_utri_c_35 * p_coeff_9 + + lsq_rmat_utri_c_36 * p_coeff_10 + ) + ) - p_coeff_2 = lsq_rmat_rdiag_c_1 * (lsq_qtmat_c_1(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_qtmat_c_1(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_qtmat_c_1(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_qtmat_c_1(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_qtmat_c_1(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_qtmat_c_1(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_qtmat_c_1(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_qtmat_c_1(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_qtmat_c_1(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - (lsq_rmat_utri_c_29*p_coeff_3 + lsq_rmat_utri_c_30*p_coeff_4 + lsq_rmat_utri_c_31*p_coeff_5 + lsq_rmat_utri_c_32*p_coeff_6 + lsq_rmat_utri_c_33*p_coeff_7 + lsq_rmat_utri_c_34*p_coeff_8 + lsq_rmat_utri_c_35*p_coeff_9 + lsq_rmat_utri_c_36*p_coeff_10)) + p_coeff_1 = p_cc - ( + p_coeff_2 * lsq_moments_1 + + p_coeff_3 * lsq_moments_2 + + p_coeff_4 * lsq_moments_3 + + p_coeff_5 * lsq_moments_4 + + p_coeff_6 * lsq_moments_5 + + p_coeff_7 * lsq_moments_6 + + p_coeff_8 * lsq_moments_7 + + p_coeff_9 * lsq_moments_8 + + p_coeff_10 * lsq_moments_9 + ) + return ( + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + ) - p_coeff_1 = p_cc - (p_coeff_2*lsq_moments_1 + p_coeff_3*lsq_moments_2 + p_coeff_4*lsq_moments_3 + p_coeff_5*lsq_moments_4 + p_coeff_6*lsq_moments_5 + p_coeff_7*lsq_moments_6 + p_coeff_8*lsq_moments_7 + p_coeff_9*lsq_moments_8 + p_coeff_10*lsq_moments_9) - return p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10 @program def recon_lsq_cell_c_stencil( @@ -138,15 +312,15 @@ def recon_lsq_cell_c_stencil( lsq_qtmat_c_7: Field[[CECECDim], float], lsq_qtmat_c_8: Field[[CECECDim], float], lsq_qtmat_c_9: Field[[CECECDim], float], - lsq_rmat_rdiag_c_1: Field[[CellDim], float], + lsq_rmat_rdiag_c_1: Field[[CellDim], float], lsq_rmat_rdiag_c_2: Field[[CellDim], float], - lsq_rmat_rdiag_c_3: Field[[CellDim], float], - lsq_rmat_rdiag_c_4: Field[[CellDim], float], - lsq_rmat_rdiag_c_5: Field[[CellDim], float], - lsq_rmat_rdiag_c_6: Field[[CellDim], float], - lsq_rmat_rdiag_c_7: Field[[CellDim], float], - lsq_rmat_rdiag_c_8: Field[[CellDim], float], - lsq_rmat_rdiag_c_9: Field[[CellDim], float], + lsq_rmat_rdiag_c_3: Field[[CellDim], float], + lsq_rmat_rdiag_c_4: Field[[CellDim], float], + lsq_rmat_rdiag_c_5: Field[[CellDim], float], + lsq_rmat_rdiag_c_6: Field[[CellDim], float], + lsq_rmat_rdiag_c_7: Field[[CellDim], float], + lsq_rmat_rdiag_c_8: Field[[CellDim], float], + lsq_rmat_rdiag_c_9: Field[[CellDim], float], lsq_rmat_utri_c_1: Field[[CellDim], float], lsq_rmat_utri_c_2: Field[[CellDim], float], lsq_rmat_utri_c_3: Field[[CellDim], float], @@ -203,5 +377,81 @@ def recon_lsq_cell_c_stencil( p_coeff_9: Field[[CellDim, KDim], float], p_coeff_10: Field[[CellDim, KDim], float], ): - _recon_lsq_cell_c_stencil(p_cc, lsq_qtmat_c_1, lsq_qtmat_c_2, lsq_qtmat_c_3, lsq_qtmat_c_4, lsq_qtmat_c_5, lsq_qtmat_c_6, lsq_qtmat_c_7, lsq_qtmat_c_8, lsq_qtmat_c_9, lsq_rmat_rdiag_c_1, lsq_rmat_rdiag_c_2, lsq_rmat_rdiag_c_3, lsq_rmat_rdiag_c_4, lsq_rmat_rdiag_c_5, lsq_rmat_rdiag_c_6, lsq_rmat_rdiag_c_7, lsq_rmat_rdiag_c_8, lsq_rmat_rdiag_c_9, lsq_rmat_utri_c_1, lsq_rmat_utri_c_2, lsq_rmat_utri_c_3, lsq_rmat_utri_c_4, lsq_rmat_utri_c_5, lsq_rmat_utri_c_6, lsq_rmat_utri_c_7, lsq_rmat_utri_c_8, lsq_rmat_utri_c_9, lsq_rmat_utri_c_10, lsq_rmat_utri_c_11, lsq_rmat_utri_c_12, lsq_rmat_utri_c_13, lsq_rmat_utri_c_14, lsq_rmat_utri_c_15, lsq_rmat_utri_c_16, lsq_rmat_utri_c_17, lsq_rmat_utri_c_18, lsq_rmat_utri_c_19, lsq_rmat_utri_c_20, lsq_rmat_utri_c_21, lsq_rmat_utri_c_22, lsq_rmat_utri_c_23, lsq_rmat_utri_c_24, lsq_rmat_utri_c_25, lsq_rmat_utri_c_26, lsq_rmat_utri_c_27, lsq_rmat_utri_c_28, lsq_rmat_utri_c_29, lsq_rmat_utri_c_30, lsq_rmat_utri_c_31, lsq_rmat_utri_c_32, lsq_rmat_utri_c_33, lsq_rmat_utri_c_34, lsq_rmat_utri_c_35, lsq_rmat_utri_c_36, lsq_moments_1, lsq_moments_2, lsq_moments_3, lsq_moments_4, lsq_moments_5, lsq_moments_6, lsq_moments_7, lsq_moments_8, lsq_moments_9, out=(p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10) - ) + _recon_lsq_cell_c_stencil( + p_cc, + lsq_qtmat_c_1, + lsq_qtmat_c_2, + lsq_qtmat_c_3, + lsq_qtmat_c_4, + lsq_qtmat_c_5, + lsq_qtmat_c_6, + lsq_qtmat_c_7, + lsq_qtmat_c_8, + lsq_qtmat_c_9, + lsq_rmat_rdiag_c_1, + lsq_rmat_rdiag_c_2, + lsq_rmat_rdiag_c_3, + lsq_rmat_rdiag_c_4, + lsq_rmat_rdiag_c_5, + lsq_rmat_rdiag_c_6, + lsq_rmat_rdiag_c_7, + lsq_rmat_rdiag_c_8, + lsq_rmat_rdiag_c_9, + lsq_rmat_utri_c_1, + lsq_rmat_utri_c_2, + lsq_rmat_utri_c_3, + lsq_rmat_utri_c_4, + lsq_rmat_utri_c_5, + lsq_rmat_utri_c_6, + lsq_rmat_utri_c_7, + lsq_rmat_utri_c_8, + lsq_rmat_utri_c_9, + lsq_rmat_utri_c_10, + lsq_rmat_utri_c_11, + lsq_rmat_utri_c_12, + lsq_rmat_utri_c_13, + lsq_rmat_utri_c_14, + lsq_rmat_utri_c_15, + lsq_rmat_utri_c_16, + lsq_rmat_utri_c_17, + lsq_rmat_utri_c_18, + lsq_rmat_utri_c_19, + lsq_rmat_utri_c_20, + lsq_rmat_utri_c_21, + lsq_rmat_utri_c_22, + lsq_rmat_utri_c_23, + lsq_rmat_utri_c_24, + lsq_rmat_utri_c_25, + lsq_rmat_utri_c_26, + lsq_rmat_utri_c_27, + lsq_rmat_utri_c_28, + lsq_rmat_utri_c_29, + lsq_rmat_utri_c_30, + lsq_rmat_utri_c_31, + lsq_rmat_utri_c_32, + lsq_rmat_utri_c_33, + lsq_rmat_utri_c_34, + lsq_rmat_utri_c_35, + lsq_rmat_utri_c_36, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + out=( + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + ), + ) diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py index 5728c519b3..870ff50b81 100644 --- a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py +++ b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py @@ -14,7 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, C2E2C2E2C, CECECDim, CellDim, KDim +from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim @field_operator @@ -38,21 +38,143 @@ def _recon_lsq_cell_c_svd_stencil( lsq_moments_7: Field[[CellDim], float], lsq_moments_8: Field[[CellDim], float], lsq_moments_9: Field[[CellDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: - p_coeff_10 = lsq_pseudoinv_9(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_9(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_9(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_9(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_9(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_9(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_9(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_9(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_9(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_9 = lsq_pseudoinv_8(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_8(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_8(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_8(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_8(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_8(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_8(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_8(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_8(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_8 = lsq_pseudoinv_7(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_7(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_7(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_7(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_7(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_7(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_7(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_7(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_7(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_7 = lsq_pseudoinv_6(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_6(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_6(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_6(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_6(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_6(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_6(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_6(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_6(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_6 = lsq_pseudoinv_5(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_5(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_5(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_5(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_5(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_5(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_5(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_5(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_5(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_5 = lsq_pseudoinv_4(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_4(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_4(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_4(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_4(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_4(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_4(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_4(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_4(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_4 = lsq_pseudoinv_3(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_3(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_3(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_3(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_3(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_3(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_3(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_3(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_3(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_3 = lsq_pseudoinv_2(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_2(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_2(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_2(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_2(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_2(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_2(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_2(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_2(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) - p_coeff_2 = lsq_pseudoinv_1(C2CECEC[0])*(p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_1(C2CECEC[1])*(p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_1(C2CECEC[2])*(p_cc(C2E2C2E2C[2]) - p_cc) + lsq_pseudoinv_1(C2CECEC[3])*(p_cc(C2E2C2E2C[3]) - p_cc) + lsq_pseudoinv_1(C2CECEC[4])*(p_cc(C2E2C2E2C[4]) - p_cc) + lsq_pseudoinv_1(C2CECEC[5])*(p_cc(C2E2C2E2C[5]) - p_cc) + lsq_pseudoinv_1(C2CECEC[6])*(p_cc(C2E2C2E2C[6]) - p_cc) + lsq_pseudoinv_1(C2CECEC[7])*(p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_1(C2CECEC[8])*(p_cc(C2E2C2E2C[8]) - p_cc) + p_coeff_10 = ( + lsq_pseudoinv_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_9 = ( + lsq_pseudoinv_8(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_8 = ( + lsq_pseudoinv_7(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_7 = ( + lsq_pseudoinv_6(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_6 = ( + lsq_pseudoinv_5(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_5 = ( + lsq_pseudoinv_4(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_4 = ( + lsq_pseudoinv_3(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_3 = ( + lsq_pseudoinv_2(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_2 = ( + lsq_pseudoinv_1(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) - p_coeff_1 = p_cc - (p_coeff_2*lsq_moments_1 + p_coeff_3*lsq_moments_2 + p_coeff_4*lsq_moments_3 + p_coeff_5*lsq_moments_4 + p_coeff_6*lsq_moments_5 + p_coeff_7*lsq_moments_6 + p_coeff_8*lsq_moments_7 + p_coeff_9*lsq_moments_8 + p_coeff_10*lsq_moments_9) + p_coeff_1 = p_cc - ( + p_coeff_2 * lsq_moments_1 + + p_coeff_3 * lsq_moments_2 + + p_coeff_4 * lsq_moments_3 + + p_coeff_5 * lsq_moments_4 + + p_coeff_6 * lsq_moments_5 + + p_coeff_7 * lsq_moments_6 + + p_coeff_8 * lsq_moments_7 + + p_coeff_9 * lsq_moments_8 + + p_coeff_10 * lsq_moments_9 + ) - return p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10 + return ( + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + ) @program @@ -88,4 +210,35 @@ def recon_lsq_cell_c_svd_stencil( p_coeff_10: Field[[CellDim, KDim], float], ): _recon_lsq_cell_c_svd_stencil( - p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, lsq_pseudoinv_3, lsq_pseudoinv_4, lsq_pseudoinv_5, lsq_pseudoinv_6, lsq_pseudoinv_7, lsq_pseudoinv_8, lsq_pseudoinv_9, lsq_moments_1, lsq_moments_2, lsq_moments_3, lsq_moments_4, lsq_moments_5, lsq_moments_6, lsq_moments_7, lsq_moments_8, lsq_moments_9, out=(p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, p_coeff_5, p_coeff_6, p_coeff_7, p_coeff_8, p_coeff_9, p_coeff_10)) + p_cc, + lsq_pseudoinv_1, + lsq_pseudoinv_2, + lsq_pseudoinv_3, + lsq_pseudoinv_4, + lsq_pseudoinv_5, + lsq_pseudoinv_6, + lsq_pseudoinv_7, + lsq_pseudoinv_8, + lsq_pseudoinv_9, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + out=( + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + ), + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py index c3971b8778..fb0c84d3f2 100644 --- a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py @@ -22,7 +22,11 @@ def _upwind_hflux_miura_stencil_02( p_cc: Field[[CellDim, KDim], float], lsq_pseudoinv_1: Field[[CECDim], float], lsq_pseudoinv_2: Field[[CECDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: p_coeff_1 = p_cc p_coeff_2 = ( lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 32548af981..5d7c602705 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -12,14 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.common.dimension import CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field, _shape -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded - +from icon4py.testutils.utils import _shape, random_field, zero_field def face_val_ppm_stencil_01_numpy( @@ -29,34 +28,60 @@ def face_val_ppm_stencil_01_numpy( elev: int32, ): - # vert_idx = np.broadcast_to(vert_idx, p_cc.shape) + # this is a comment: vert_idx = np.broadcast_to(vert_idx, p_cc.shape) # 01a - zfac_m1 = (p_cc[:,1:-1] - p_cc[:,:-2]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,:-2]) - zfac = (p_cc[:,2:] - p_cc[:,1:-1]) / (p_cellhgt_mc_now[:,2:] + p_cellhgt_mc_now[:,1:-1]) - z_slope_a = ( p_cellhgt_mc_now[:,1:-1] / (p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,2:]) ) * ( (2.*p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1]) * zfac + (p_cellhgt_mc_now[:,1:-1] + 2.*p_cellhgt_mc_now[:,2:]) * zfac_m1) - + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / ( + p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1] + ) + z_slope_a = ( + p_cellhgt_mc_now[:, 1:-1] + / ( + p_cellhgt_mc_now[:, :-2] + + p_cellhgt_mc_now[:, 1:-1] + + p_cellhgt_mc_now[:, 2:] + ) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 + ) # 01b - zfac_m1 = (p_cc[:,1:-1] - p_cc[:,:-2]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,:-2]) - zfac = (p_cc[:,1:-1] - p_cc[:,1:-1]) / (p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,1:-1]) - z_slope_b = ( p_cellhgt_mc_now[:,1:-1] / (p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1] + p_cellhgt_mc_now[:,1:-1]) ) * ( (2.*p_cellhgt_mc_now[:,:-2] + p_cellhgt_mc_now[:,1:-1]) * zfac + (p_cellhgt_mc_now[:,1:-1] + 2.*p_cellhgt_mc_now[:,1:-1]) * zfac_m1) + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1] + ) + z_slope_b = ( + p_cellhgt_mc_now[:, 1:-1] + / ( + p_cellhgt_mc_now[:, :-2] + + p_cellhgt_mc_now[:, 1:-1] + + p_cellhgt_mc_now[:, 1:-1] + ) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 + ) - z_slope = np.where(vert_idx[1:-1] np.ndarray: tmp = pd_time * ( - p_mflx_contra_v[:,1:] * deepatmo_divzl - p_mflx_contra_v[:,:-1] * deepatmo_divzu + p_mflx_contra_v[:, 1:] * deepatmo_divzl + - p_mflx_contra_v[:, :-1] * deepatmo_divzu ) return rhodz_ast + tmp diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py index 2aff78df0e..d944672704 100644 --- a/advection/tests/test_step_advection_stencil_02.py +++ b/advection/tests/test_step_advection_stencil_02.py @@ -28,8 +28,8 @@ def step_advection_stencil_02_numpy( ) -> np.ndarray: tmp = ( - p_mflx_contra_v[:,1:] * deepatmo_divzl - - p_mflx_contra_v[:,:-1] * deepatmo_divzu + p_mflx_contra_v[:, 1:] * deepatmo_divzl + - p_mflx_contra_v[:, :-1] * deepatmo_divzu ) return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py index 4f034feed8..8722fc004f 100644 --- a/advection/tests/test_vert_adv_stencil_01.py +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -32,7 +32,10 @@ def vert_adv_stencil_01_numpy( tracer_new = ( tracer_now * rhodz_now + p_dtime - * (p_mflx_tracer_v[:,1:] * deepatmo_divzl - p_mflx_tracer_v[:,:-1] * deepatmo_divzu) + * ( + p_mflx_tracer_v[:, 1:] * deepatmo_divzl + - p_mflx_tracer_v[:, :-1] * deepatmo_divzu + ) ) / rhodz_new return tracer_new diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 2765918b86..6bcaadb4f7 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -26,21 +26,21 @@ def v_limit_prbl_sm_stencil_01_numpy( p_cc: np.array, ): - z_delta = p_face[:,:-1] - p_face[:,1:] - z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:,:-1] + p_face[:,1:])) + z_delta = p_face[:, :-1] - p_face[:, 1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) q_face_up, q_face_low = np.where( np.abs(z_delta) < -1 * z_a6i, np.where( - (p_cc < np.minimum(p_face[:,:-1], p_face[:,1:])), + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), (p_cc, p_cc), np.where( - p_face[:,:-1] > p_face[:,1:], - (3.0 * p_cc - 2.0 * p_face[:,1:], p_face[:,1:]), - (p_face[:,:-1], 3.0 * p_cc - 2.0 * p_face[:,:-1]), + p_face[:, :-1] > p_face[:, 1:], + (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), + (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), ), ), - (p_face[:,:-1], p_face[:,1:]), + (p_face[:, :-1], p_face[:, 1:]), ) return q_face_up, q_face_low diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index 33f72bae2d..c2c322f55d 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -365,6 +365,7 @@ class ImportsStatementGenerator(TemplatedGenerator): class StartCreateStatement(eve.Node): extra_fields: Optional[list[str]] + class StartCreateStatementGenerator(TemplatedGenerator): StartCreateStatement = as_jinja( """ diff --git a/liskov/src/icon4py/liskov/codegen/generate.py b/liskov/src/icon4py/liskov/codegen/generate.py index 83fa0c1b41..ee622e36f4 100644 --- a/liskov/src/icon4py/liskov/codegen/generate.py +++ b/liskov/src/icon4py/liskov/codegen/generate.py @@ -205,7 +205,7 @@ def _generate_end_stencil(self) -> None: profile=self.profile, noendif=self.directives.EndStencil[i].noendif, noprofile=self.directives.EndStencil[i].noprofile, - noaccenddata=self.directives.EndStencil[i].noaccenddata + noaccenddata=self.directives.EndStencil[i].noaccenddata, ) def _generate_imports(self) -> None: @@ -222,24 +222,24 @@ def _generate_imports(self) -> None: def _generate_create(self) -> None: """Generate f90 code for OpenACC DATA CREATE statements.""" if self.directives.StartCreate != UnusedDirective: - for startcreate in self.directives.StartCreate: # type: ignore - logger.info("Generating DATA CREATE statement.") - self._generate( - StartCreateStatement, - StartCreateStatementGenerator, - startcreate.startln, - startcreate.endln, - extra_fields=startcreate.extra_fields, - ) + for startcreate in self.directives.StartCreate: # type: ignore + logger.info("Generating DATA CREATE statement.") + self._generate( + StartCreateStatement, + StartCreateStatementGenerator, + startcreate.startln, + startcreate.endln, + extra_fields=startcreate.extra_fields, + ) if self.directives.EndCreate != UnusedDirective: - for endcreate in self.directives.EndCreate: # type: ignore - self._generate( - EndCreateStatement, - EndCreateStatementGenerator, - endcreate.startln, - endcreate.endln, - ) + for endcreate in self.directives.EndCreate: # type: ignore + self._generate( + EndCreateStatement, + EndCreateStatementGenerator, + endcreate.startln, + endcreate.endln, + ) def _generate_endif(self) -> None: """Generate f90 code for endif statements.""" diff --git a/liskov/src/icon4py/liskov/parsing/deserialise.py b/liskov/src/icon4py/liskov/parsing/deserialise.py index 8bb01ac928..aad10f8f60 100644 --- a/liskov/src/icon4py/liskov/parsing/deserialise.py +++ b/liskov/src/icon4py/liskov/parsing/deserialise.py @@ -167,7 +167,11 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData]: extra_fields = named_args["extra_fields"].split(",") deserialised.append( - self.dtype( startln=directive.startln, endln=directive.endln, extra_fields=extra_fields ) + self.dtype( + startln=directive.startln, + endln=directive.endln, + extra_fields=extra_fields, + ) ) return deserialised diff --git a/liskov/tests/test_generation.py b/liskov/tests/test_generation.py index 6b4a6d8926..7c99e1528f 100644 --- a/liskov/tests/test_generation.py +++ b/liskov/tests/test_generation.py @@ -71,7 +71,12 @@ def serialised_directives(): copies=True, ) end_stencil_data = EndStencilData( - name="stencil1", startln=3, endln=4, noendif=False, noprofile=False, noaccenddata=False + name="stencil1", + startln=3, + endln=4, + noendif=False, + noprofile=False, + noaccenddata=False, ) declare_data = DeclareData( startln=5, diff --git a/liskov/tests/test_validation.py b/liskov/tests/test_validation.py index 885c70d3bb..9cafa0507b 100644 --- a/liskov/tests/test_validation.py +++ b/liskov/tests/test_validation.py @@ -23,12 +23,7 @@ UnbalancedStencilDirectiveError, ) from icon4py.liskov.parsing.parse import DirectivesParser -from icon4py.liskov.parsing.types import ( - Declare, - Imports, - StartCreate, - StartStencil, -) +from icon4py.liskov.parsing.types import Declare, Imports, StartStencil from icon4py.liskov.parsing.validation import DirectiveSyntaxValidator From cf990a1283f34267ead75c596ca6e6e50ce6573b Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 8 Jun 2023 11:36:55 +0200 Subject: [PATCH 059/105] update style --- tools/src/icon4pytools/f2ser/cli.py | 1 - tools/src/icon4pytools/f2ser/parse.py | 3 +-- tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py | 1 - tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py | 1 - tools/src/icon4pytools/icon4pygen/bindings/locations.py | 4 ++-- tools/src/icon4pytools/icon4pygen/bindings/utils.py | 1 - tools/src/icon4pytools/icon4pygen/icochainsize.py | 1 + tools/src/icon4pytools/icon4pygen/metadata.py | 4 ++-- .../src/icon4pytools/liskov/codegen/integration/template.py | 1 - tools/src/icon4pytools/liskov/external/gt4py.py | 1 - tools/src/icon4pytools/liskov/external/metadata.py | 1 - tools/tests/f2ser/test_f2ser_cli.py | 1 - tools/tests/f2ser/test_f2ser_codegen.py | 1 - tools/tests/f2ser/test_granule_deserialiser.py | 1 - tools/tests/f2ser/test_parsing.py | 1 - tools/tests/icon4pygen/test_backend.py | 1 - tools/tests/icon4pygen/test_codegen.py | 4 ++-- tools/tests/icon4pygen/test_field_rendering.py | 4 ++-- tools/tests/icon4pygen/test_icochainsize.py | 1 - tools/tests/liskov/conftest.py | 3 +-- tools/tests/liskov/test_cli.py | 1 - tools/tests/liskov/test_directives_deserialiser.py | 3 +-- tools/tests/liskov/test_generation.py | 1 - tools/tests/liskov/test_parser.py | 5 ++--- tools/tests/liskov/test_scanner.py | 3 +-- tools/tests/liskov/test_serialisation_deserialiser.py | 1 - tools/tests/liskov/test_utils.py | 3 +-- tools/tests/liskov/test_validation.py | 3 +-- 28 files changed, 17 insertions(+), 39 deletions(-) diff --git a/tools/src/icon4pytools/f2ser/cli.py b/tools/src/icon4pytools/f2ser/cli.py index 290d8a54f9..0fe794a5b1 100644 --- a/tools/src/icon4pytools/f2ser/cli.py +++ b/tools/src/icon4pytools/f2ser/cli.py @@ -15,7 +15,6 @@ from typing import Optional import click - from icon4pytools.f2ser.deserialise import ParsedGranuleDeserialiser from icon4pytools.f2ser.parse import GranuleParser from icon4pytools.liskov.codegen.serialisation.generate import SerialisationCodeGenerator diff --git a/tools/src/icon4pytools/f2ser/parse.py b/tools/src/icon4pytools/f2ser/parse.py index 28fdacb725..c4addb3870 100644 --- a/tools/src/icon4pytools/f2ser/parse.py +++ b/tools/src/icon4pytools/f2ser/parse.py @@ -17,9 +17,8 @@ from pathlib import Path from typing import Any, Optional -from numpy.f2py.crackfortran import crackfortran - from icon4pytools.f2ser.exceptions import MissingDerivedTypeError, ParsingError +from numpy.f2py.crackfortran import crackfortran def crack(path: Path) -> dict: diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py index ba6951f572..b31729f887 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py @@ -17,7 +17,6 @@ from gt4py import eve from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import Node, TemplatedGenerator, format_source - from icon4pytools.icon4pygen.bindings.codegen.header import ( CppFreeFunc, CppRunAndVerifyFuncDeclaration, diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py index 87ffad337c..671e755467 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py @@ -16,7 +16,6 @@ from typing import Union from gt4py.next.type_system import type_specifications as ts - from icon4pytools.icon4pygen.bindings.locations import ( BasicLocation, ChainedLocation, diff --git a/tools/src/icon4pytools/icon4pygen/bindings/locations.py b/tools/src/icon4pytools/icon4pygen/bindings/locations.py index 4071da48b6..5a35054e75 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/locations.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/locations.py @@ -15,10 +15,10 @@ from typing import Iterator from gt4py.next.ffront.fbuiltins import Dimension -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim - from icon4pytools.icon4pygen.bindings.codegen.render.location import LocationRenderer +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim + class BasicLocation: renderer = LocationRenderer diff --git a/tools/src/icon4pytools/icon4pygen/bindings/utils.py b/tools/src/icon4pytools/icon4pygen/bindings/utils.py index 5d5cb72c1e..c372266237 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/utils.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/utils.py @@ -17,7 +17,6 @@ from pathlib import Path from gt4py.next.common import Dimension - from icon4pytools.icon4pygen.icochainsize import IcoChainSize diff --git a/tools/src/icon4pytools/icon4pygen/icochainsize.py b/tools/src/icon4pytools/icon4pygen/icochainsize.py index 38f3b825ac..66b220a1cc 100644 --- a/tools/src/icon4pytools/icon4pygen/icochainsize.py +++ b/tools/src/icon4pytools/icon4pygen/icochainsize.py @@ -111,6 +111,7 @@ from typing import List, TypeAlias from gt4py.next.common import Dimension + from icon4py.common.dimension import CellDim, EdgeDim, VertexDim diff --git a/tools/src/icon4pytools/icon4pygen/metadata.py b/tools/src/icon4pytools/icon4pygen/metadata.py index 1f1b42e4cb..340f811993 100644 --- a/tools/src/icon4pytools/icon4pygen/metadata.py +++ b/tools/src/icon4pytools/icon4pygen/metadata.py @@ -25,14 +25,14 @@ from gt4py.next.iterator import ir as itir from gt4py.next.iterator.runtime import FendefDispatcher from gt4py.next.type_system import type_specifications as ts -from icon4py.common.dimension import CellDim, EdgeDim, Koff, VertexDim - from icon4pytools.icon4pygen.exceptions import ( InvalidConnectivityException, MultipleFieldOperatorException, ) from icon4pytools.icon4pygen.icochainsize import IcoChainSize +from icon4py.common.dimension import CellDim, EdgeDim, Koff, VertexDim + @dataclass(frozen=True) class StencilInfo: diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index dba9543864..afa5a7406f 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -18,7 +18,6 @@ import gt4py.eve as eve from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import TemplatedGenerator - from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError from icon4pytools.liskov.codegen.integration.interface import DeclareData, StartStencilData from icon4pytools.liskov.external.metadata import CodeMetadata diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 4ad14ef9d3..6996628825 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -16,7 +16,6 @@ from typing import Any from gt4py.next.ffront.decorator import Program - from icon4pytools.common.logger import setup_logger from icon4pytools.icon4pygen.metadata import get_stencil_info from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface diff --git a/tools/src/icon4pytools/liskov/external/metadata.py b/tools/src/icon4pytools/liskov/external/metadata.py index b9d6728fdf..486fd92059 100644 --- a/tools/src/icon4pytools/liskov/external/metadata.py +++ b/tools/src/icon4pytools/liskov/external/metadata.py @@ -14,7 +14,6 @@ from typing import Any import click - from icon4pytools import __version__ from icon4pytools.liskov.external.exceptions import MissingClickContextError diff --git a/tools/tests/f2ser/test_f2ser_cli.py b/tools/tests/f2ser/test_f2ser_cli.py index 83c4ea1691..64d36c8392 100644 --- a/tools/tests/f2ser/test_f2ser_cli.py +++ b/tools/tests/f2ser/test_f2ser_cli.py @@ -13,7 +13,6 @@ import pytest from click.testing import CliRunner - from icon4pytools.f2ser.cli import main from icon4pytools.f2ser.exceptions import MissingDerivedTypeError diff --git a/tools/tests/f2ser/test_f2ser_codegen.py b/tools/tests/f2ser/test_f2ser_codegen.py index 5ca607eff9..7efd6793b6 100644 --- a/tools/tests/f2ser/test_f2ser_codegen.py +++ b/tools/tests/f2ser/test_f2ser_codegen.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.f2ser.deserialise import ParsedGranuleDeserialiser from icon4pytools.f2ser.parse import GranuleParser from icon4pytools.liskov.codegen.serialisation.generate import SerialisationCodeGenerator diff --git a/tools/tests/f2ser/test_granule_deserialiser.py b/tools/tests/f2ser/test_granule_deserialiser.py index d0f478e7ce..d83880ae98 100644 --- a/tools/tests/f2ser/test_granule_deserialiser.py +++ b/tools/tests/f2ser/test_granule_deserialiser.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.f2ser.deserialise import ParsedGranuleDeserialiser from icon4pytools.f2ser.parse import CodegenContext, GranuleParser, ParsedGranule from icon4pytools.liskov.codegen.serialisation.interface import ( diff --git a/tools/tests/f2ser/test_parsing.py b/tools/tests/f2ser/test_parsing.py index 177fc496ee..cf06d81e00 100644 --- a/tools/tests/f2ser/test_parsing.py +++ b/tools/tests/f2ser/test_parsing.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.f2ser.exceptions import MissingDerivedTypeError, ParsingError from icon4pytools.f2ser.parse import CodegenContext, GranuleParser diff --git a/tools/tests/icon4pygen/test_backend.py b/tools/tests/icon4pygen/test_backend.py index e712d3ffa6..fed703c73a 100644 --- a/tools/tests/icon4pygen/test_backend.py +++ b/tools/tests/icon4pygen/test_backend.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest from gt4py.next.iterator import ir as itir - from icon4pytools.icon4pygen import backend from icon4pytools.icon4pygen.backend import GTHeader diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 0bb9ad5c2c..de476bb761 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -15,12 +15,12 @@ import pkgutil import re -import icon4py.atm_dyn_iconam import pytest from click.testing import CliRunner - from icon4pytools.icon4pygen.cli import main +import icon4py.atm_dyn_iconam + from .helpers import get_stencil_module_path diff --git a/tools/tests/icon4pygen/test_field_rendering.py b/tools/tests/icon4pygen/test_field_rendering.py index 2ace9a6e24..f6649b7c62 100644 --- a/tools/tests/icon4pygen/test_field_rendering.py +++ b/tools/tests/icon4pygen/test_field_rendering.py @@ -13,11 +13,11 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2CDim, EdgeDim, KDim - from icon4pytools.icon4pygen.bindings.workflow import PyBindGen from icon4pytools.icon4pygen.metadata import get_stencil_info +from icon4py.common.dimension import E2CDim, EdgeDim, KDim + def test_horizontal_field_sid_rendering(): @field_operator diff --git a/tools/tests/icon4pygen/test_icochainsize.py b/tools/tests/icon4pygen/test_icochainsize.py index f9e73c5296..f606153fc5 100644 --- a/tools/tests/icon4pygen/test_icochainsize.py +++ b/tools/tests/icon4pygen/test_icochainsize.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.icon4pygen.exceptions import InvalidConnectivityException from icon4pytools.icon4pygen.metadata import provide_offset diff --git a/tools/tests/liskov/conftest.py b/tools/tests/liskov/conftest.py index 416941b8d7..b67bba7834 100644 --- a/tools/tests/liskov/conftest.py +++ b/tools/tests/liskov/conftest.py @@ -13,10 +13,9 @@ from pathlib import Path +import icon4pytools.liskov.parsing.parse as ts import pytest from click.testing import CliRunner - -import icon4pytools.liskov.parsing.parse as ts from icon4pytools.liskov.parsing.scan import DirectivesScanner diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index edd178f62e..ca76624e93 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -14,7 +14,6 @@ import itertools import pytest - from icon4pytools.liskov.cli import main from icon4pytools.liskov.external.exceptions import MissingCommandError diff --git a/tools/tests/liskov/test_directives_deserialiser.py b/tools/tests/liskov/test_directives_deserialiser.py index 13a235f6f5..55d5aba19b 100644 --- a/tools/tests/liskov/test_directives_deserialiser.py +++ b/tools/tests/liskov/test_directives_deserialiser.py @@ -13,9 +13,8 @@ import unittest -import pytest - import icon4pytools.liskov.parsing.parse as ts +import pytest from icon4pytools.liskov.codegen.integration.deserialise import ( DeclareDataFactory, EndCreateDataFactory, diff --git a/tools/tests/liskov/test_generation.py b/tools/tests/liskov/test_generation.py index dc16b09c6a..f33184a0b9 100644 --- a/tools/tests/liskov/test_generation.py +++ b/tools/tests/liskov/test_generation.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.liskov.codegen.integration.generate import IntegrationCodeGenerator from icon4pytools.liskov.codegen.integration.interface import ( BoundsData, diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index d0d8aa17cb..0d6cb1d97d 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -14,12 +14,11 @@ from collections import defaultdict -import pytest -from pytest import mark - import icon4pytools.liskov.parsing.parse as ts +import pytest from icon4pytools.liskov.parsing.exceptions import UnsupportedDirectiveError from icon4pytools.liskov.parsing.parse import DirectivesParser +from pytest import mark from .conftest import insert_new_lines, scan_for_directives from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL diff --git a/tools/tests/liskov/test_scanner.py b/tools/tests/liskov/test_scanner.py index 14d2839f45..31b6a8c519 100644 --- a/tools/tests/liskov/test_scanner.py +++ b/tools/tests/liskov/test_scanner.py @@ -15,11 +15,10 @@ from pathlib import Path import pytest -from pytest import mark - from icon4pytools.liskov.parsing.exceptions import DirectiveSyntaxError from icon4pytools.liskov.parsing.scan import DirectivesScanner from icon4pytools.liskov.parsing.types import RawDirective +from pytest import mark from .fortran_samples import DIRECTIVES_SAMPLE, NO_DIRECTIVES_STENCIL diff --git a/tools/tests/liskov/test_serialisation_deserialiser.py b/tools/tests/liskov/test_serialisation_deserialiser.py index 0431086beb..755f6955cf 100644 --- a/tools/tests/liskov/test_serialisation_deserialiser.py +++ b/tools/tests/liskov/test_serialisation_deserialiser.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest - from icon4pytools.liskov.codegen.serialisation.deserialise import ( InitDataFactory, SavepointDataFactory, diff --git a/tools/tests/liskov/test_utils.py b/tools/tests/liskov/test_utils.py index 82d5a742e4..a642de6647 100644 --- a/tools/tests/liskov/test_utils.py +++ b/tools/tests/liskov/test_utils.py @@ -12,9 +12,8 @@ # SPDX-License-Identifier: GPL-3.0-or-later from copy import deepcopy -import pytest - import icon4pytools.liskov.parsing.parse as ts +import pytest from icon4pytools.liskov.parsing.utils import ( extract_directive, print_parsed_directive, diff --git a/tools/tests/liskov/test_validation.py b/tools/tests/liskov/test_validation.py index d6fbc4f927..6ef225aed2 100644 --- a/tools/tests/liskov/test_validation.py +++ b/tools/tests/liskov/test_validation.py @@ -12,8 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest -from pytest import mark - from icon4pytools.liskov.parsing.exceptions import ( DirectiveSyntaxError, RepeatedDirectiveError, @@ -22,6 +20,7 @@ ) from icon4pytools.liskov.parsing.parse import Declare, DirectivesParser, Imports, StartStencil from icon4pytools.liskov.parsing.validation import DirectiveSyntaxValidator +from pytest import mark from .conftest import insert_new_lines, scan_for_directives from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL From 11f805b965d0fa91a40d2ffd54c5fc84c8c20197 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 12 Jun 2023 09:57:35 +0200 Subject: [PATCH 060/105] split stencil --- .../advection/v_limit_prbl_sm_stencil_01.py | 29 +++------ .../advection/v_limit_prbl_sm_stencil_02.py | 61 +++++++++++++++++++ 2 files changed, 69 insertions(+), 21 deletions(-) create mode 100644 advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py index 3514ff86bb..d4101936a5 100644 --- a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.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, FieldOffset, abs, minimum, where +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, abs, minimum, where, int32 from icon4py.common.dimension import CellDim, KDim @@ -24,37 +24,24 @@ def _v_limit_prbl_sm_stencil_01( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> Field[[CellDim, KDim], int32]: z_delta = p_face - p_face(Koff[1]) - z_a6i = -6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) - - q_face_up, q_face_low = where( - abs(z_delta) < z_a6i, - where( - (p_cc < minimum(p_face, p_face(Koff[1]))), - (p_cc, p_cc), - where( - p_face > p_face(Koff[1]), - (3.0 * p_cc - 2.0 * p_face(Koff[1]), p_face(Koff[1])), - (p_face, 3.0 * p_cc - 2.0 * p_face), - ), - ), - (p_face, p_face(Koff[1])), - ) + z_a6i = 6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) - return q_face_up, q_face_low + l_limit = where(abs(z_delta) < -1.0*z_a6i, int32(1), int32(0)) + + return l_limit @program def v_limit_prbl_sm_stencil_01( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], - p_face_up: Field[[CellDim, KDim], float], - p_face_low: Field[[CellDim, KDim], float], + l_limit: Field[[CellDim, KDim], int32], ): _v_limit_prbl_sm_stencil_01( p_face, p_cc, - out=(p_face_up, p_face_low), + out=l_limit, ) diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py new file mode 100644 index 0000000000..30b9cdadea --- /dev/null +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py @@ -0,0 +1,61 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, abs, minimum, where, int32 + +from icon4py.common.dimension import CellDim, KDim + + +Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) + + +@field_operator +def _v_limit_prbl_sm_stencil_02( + l_limit: Field[[CellDim, KDim], int32], + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + + q_face_up, q_face_low = where( + l_limit!=int32(0), + where( + (p_cc < minimum(p_face, p_face(Koff[1]))), + (p_cc, p_cc), + where( + p_face > p_face(Koff[1]), + (3.0 * p_cc - 2.0 * p_face(Koff[1]), p_face(Koff[1])), + (p_face, 3.0 * p_cc - 2.0 * p_face), + ), + ), + (p_face, p_face(Koff[1])), + ) + + return q_face_up, q_face_low + + +@program +def v_limit_prbl_sm_stencil_02( + l_limit: Field[[CellDim, KDim], int32], + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + p_face_up: Field[[CellDim, KDim], float], + p_face_low: Field[[CellDim, KDim], float], +): + _v_limit_prbl_sm_stencil_02( + l_limit, + p_face, + p_cc, + out=(p_face_up, p_face_low), + ) From 577b548c8c39bed44564151932573745d94dc32b Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 12 Jun 2023 13:15:21 +0200 Subject: [PATCH 061/105] fix test utils --- .../tests/test_face_val_ppm_stencil_01.py | 4 +- .../tests/test_face_val_ppm_stencil_02.py | 4 +- .../tests/test_face_val_ppm_stencil_05.py | 4 +- .../tests/test_hflx_limiter_mo_stencil_02.py | 4 +- .../tests/test_hflx_limiter_mo_stencil_03.py | 4 +- .../tests/test_hflx_limiter_mo_stencil_04.py | 4 +- .../tests/test_hflx_limiter_pd_stencil_01.py | 4 +- .../tests/test_hflx_limiter_pd_stencil_02.py | 4 +- advection/tests/test_hor_adv_stencil_01.py | 4 +- .../tests/test_rbf_intp_edge_stencil_01.py | 4 +- advection/tests/test_set_zero_c.py | 4 +- advection/tests/test_set_zero_c_k.py | 4 +- .../tests/test_step_advection_stencil_01.py | 4 +- .../tests/test_step_advection_stencil_02.py | 4 +- .../tests/test_step_advection_stencil_03.py | 4 +- .../tests/test_step_advection_stencil_04.py | 4 +- .../test_upwind_hflux_miura_stencil_01.py | 4 +- .../test_upwind_hflux_miura_stencil_02.py | 4 +- .../tests/test_upwind_vflux_ppm_stencil_01.py | 4 +- advection/tests/test_vert_adv_stencil_01.py | 4 +- .../tests/test_vlimit_prbl_sm_stencil_01.py | 33 +++------ .../tests/test_vlimit_prbl_sm_stencil_02.py | 70 +++++++++++++++++++ atm_dyn_iconam/tests/test_compute_airmass.py | 4 +- 23 files changed, 122 insertions(+), 65 deletions(-) create mode 100644 advection/tests/test_vlimit_prbl_sm_stencil_02.py diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 5d7c602705..248dc991ac 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -17,8 +17,8 @@ from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import _shape, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import _shape, random_field, zero_field def face_val_ppm_stencil_01_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py index a437e84e44..56a82e1127 100644 --- a/advection/tests/test_face_val_ppm_stencil_02.py +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -17,8 +17,8 @@ from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import _shape, random_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import _shape, random_field def face_val_ppm_stencil_02_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/advection/tests/test_face_val_ppm_stencil_05.py index 08bd78fc46..96d17c85fc 100644 --- a/advection/tests/test_face_val_ppm_stencil_05.py +++ b/advection/tests/test_face_val_ppm_stencil_05.py @@ -15,8 +15,8 @@ from icon4py.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def face_val_ppm_stencil_05_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py index e786bd4545..89f1e9ea78 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -18,8 +18,8 @@ hflx_limiter_mo_stencil_02, ) from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import constant_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import constant_field, random_field, zero_field def hflx_limiter_mo_stencil_02_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py index d2901d96ca..c1b1c0b3e2 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -17,8 +17,8 @@ hflx_limiter_mo_stencil_03_min_max, ) from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def hflx_limiter_mo_stencil_03_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/advection/tests/test_hflx_limiter_mo_stencil_04.py index bdb7211b88..8c0436bcdd 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -17,8 +17,8 @@ hflx_limiter_mo_stencil_04, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def hflx_limiter_mo_stencil_04_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/advection/tests/test_hflx_limiter_pd_stencil_01.py index 92fa59ca90..25311d1bc3 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -18,8 +18,8 @@ hflx_limiter_pd_stencil_01, ) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import as_1D_sparse_field, random_field, zero_field def hflx_limiter_pd_stencil_01_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index 57eb18aa59..8ef646fe81 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -17,8 +17,8 @@ hflx_limiter_pd_stencil_02, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import constant_field, random_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import constant_field, random_field def hflx_limiter_pd_stencil_02_numpy( diff --git a/advection/tests/test_hor_adv_stencil_01.py b/advection/tests/test_hor_adv_stencil_01.py index 0dbbe5595c..1274e7f2fd 100644 --- a/advection/tests/test_hor_adv_stencil_01.py +++ b/advection/tests/test_hor_adv_stencil_01.py @@ -16,8 +16,8 @@ from icon4py.advection.hor_adv_stencil_01 import hor_adv_stencil_01 from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import as_1D_sparse_field, random_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import as_1D_sparse_field, random_field def hor_adv_stencil_01_numpy( diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/advection/tests/test_rbf_intp_edge_stencil_01.py index 84bd29aa43..80e0a5f3a7 100644 --- a/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -15,8 +15,8 @@ from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def rbf_intp_edge_stencil_01_numpy( diff --git a/advection/tests/test_set_zero_c.py b/advection/tests/test_set_zero_c.py index 018cc1fc3a..4d887612f0 100644 --- a/advection/tests/test_set_zero_c.py +++ b/advection/tests/test_set_zero_c.py @@ -15,8 +15,8 @@ from icon4py.advection.set_zero_c import set_zero_c from icon4py.common.dimension import CellDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def test_set_zero_cell_k(): diff --git a/advection/tests/test_set_zero_c_k.py b/advection/tests/test_set_zero_c_k.py index 1323924730..9d23cf8b6d 100644 --- a/advection/tests/test_set_zero_c_k.py +++ b/advection/tests/test_set_zero_c_k.py @@ -15,8 +15,8 @@ from icon4py.advection.set_zero_c_k import set_zero_c_k from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def test_set_zero_c_k(): diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py index 68afdc0b44..8c8d6960a0 100644 --- a/advection/tests/test_step_advection_stencil_01.py +++ b/advection/tests/test_step_advection_stencil_01.py @@ -15,8 +15,8 @@ from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def step_advection_stencil_01_numpy( diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py index d944672704..e735bc8a4a 100644 --- a/advection/tests/test_step_advection_stencil_02.py +++ b/advection/tests/test_step_advection_stencil_02.py @@ -15,8 +15,8 @@ from icon4py.advection.step_advection_stencil_02 import step_advection_stencil_02 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def step_advection_stencil_02_numpy( diff --git a/advection/tests/test_step_advection_stencil_03.py b/advection/tests/test_step_advection_stencil_03.py index ade783cb04..33997d8b11 100644 --- a/advection/tests/test_step_advection_stencil_03.py +++ b/advection/tests/test_step_advection_stencil_03.py @@ -15,8 +15,8 @@ from icon4py.advection.step_advection_stencil_03 import step_advection_stencil_03 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field def step_advection_stencil_03_numpy( diff --git a/advection/tests/test_step_advection_stencil_04.py b/advection/tests/test_step_advection_stencil_04.py index 6da3ec9b54..fd99e37083 100644 --- a/advection/tests/test_step_advection_stencil_04.py +++ b/advection/tests/test_step_advection_stencil_04.py @@ -15,8 +15,8 @@ from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def step_advection_stencil_04_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/advection/tests/test_upwind_hflux_miura_stencil_01.py index a4fa164a72..dbdd357c78 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -17,8 +17,8 @@ upwind_hflux_miura_stencil_01, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def upwind_hflux_miura_stencil_01_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py index 81bf89beb4..13ce1b5077 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_02.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -18,8 +18,8 @@ upwind_hflux_miura_stencil_02, ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import as_1D_sparse_field, random_field, zero_field def upwind_hflux_miura_stencil_02_numpy( diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 36f0d8246e..96f2990e45 100644 --- a/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -17,8 +17,8 @@ upwind_vflux_ppm_stencil_01, ) from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def upwind_vflux_ppm_stencil_01_numpy( diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py index 8722fc004f..15416c3d58 100644 --- a/advection/tests/test_vert_adv_stencil_01.py +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -15,8 +15,8 @@ from icon4py.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field def vert_adv_stencil_01_numpy( diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 6bcaadb4f7..0a20d24798 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -17,8 +17,10 @@ v_limit_prbl_sm_stencil_01, ) from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field +from gt4py.next.ffront.fbuiltins import int32 + def v_limit_prbl_sm_stencil_01_numpy( @@ -29,31 +31,18 @@ def v_limit_prbl_sm_stencil_01_numpy( z_delta = p_face[:, :-1] - p_face[:, 1:] z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - q_face_up, q_face_low = np.where( - np.abs(z_delta) < -1 * z_a6i, - np.where( - (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), - (p_cc, p_cc), - np.where( - p_face[:, :-1] > p_face[:, 1:], - (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), - (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), - ), - ), - (p_face[:, :-1], p_face[:, 1:]), - ) + l_limit = np.where( np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) - return q_face_up, q_face_low + return l_limit def test_v_limit_prbl_sm_stencil_01(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - p_face_up = zero_field(mesh, CellDim, KDim) - p_face_low = zero_field(mesh, CellDim, KDim) + l_limit = zero_field(mesh, CellDim, KDim, dtype=in32) - p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_01_numpy( + l_limit_ref = v_limit_prbl_sm_stencil_01_numpy( np.asarray(p_face), np.asarray(p_cc), ) @@ -61,10 +50,8 @@ def test_v_limit_prbl_sm_stencil_01(): v_limit_prbl_sm_stencil_01( p_face, p_cc, - p_face_up, - p_face_low, + l_limit, offset_provider={"Koff": KDim}, ) - assert np.allclose(p_face_up_ref[:, :-1], p_face_up[:, :-1]) - assert np.allclose(p_face_low_ref[:, :-1], p_face_low[:, :-1]) + assert np.allclose(l_limit_ref, l_limit) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/advection/tests/test_vlimit_prbl_sm_stencil_02.py new file mode 100644 index 0000000000..cef8b8d589 --- /dev/null +++ b/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -0,0 +1,70 @@ +# 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 + +import numpy as np + +from icon4py.advection.v_limit_prbl_sm_stencil_01 import ( + v_limit_prbl_sm_stencil_01, +) +from icon4py.common.dimension import CellDim, KDim +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.utils import random_field, zero_field + + +def v_limit_prbl_sm_stencil_01_numpy( + p_face: np.array, + p_cc: np.array, +): + + z_delta = p_face[:, :-1] - p_face[:, 1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) + + q_face_up, q_face_low = np.where( + np.abs(z_delta) < -1 * z_a6i, + np.where( + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), + (p_cc, p_cc), + np.where( + p_face[:, :-1] > p_face[:, 1:], + (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), + (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), + ), + ), + (p_face[:, :-1], p_face[:, 1:]), + ) + + return q_face_up, q_face_low + + +def test_v_limit_prbl_sm_stencil_01(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + p_face_up = zero_field(mesh, CellDim, KDim) + p_face_low = zero_field(mesh, CellDim, KDim) + + p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_01_numpy( + np.asarray(p_face), + np.asarray(p_cc), + ) + + v_limit_prbl_sm_stencil_01( + p_face, + p_cc, + p_face_up, + p_face_low, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(p_face_up_ref[:, :-1], p_face_up[:, :-1]) + assert np.allclose(p_face_low_ref[:, :-1], p_face_low[:, :-1]) diff --git a/atm_dyn_iconam/tests/test_compute_airmass.py b/atm_dyn_iconam/tests/test_compute_airmass.py index 881838f8b5..1742984fcd 100644 --- a/atm_dyn_iconam/tests/test_compute_airmass.py +++ b/atm_dyn_iconam/tests/test_compute_airmass.py @@ -15,8 +15,8 @@ from icon4py.atm_dyn_iconam.compute_airmass import compute_airmass from icon4py.common.dimension import CellDim, KDim -from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field +from .test_utils.simple_mesh import SimpleMesh +from .test_utils.helpers import random_field def compute_airmass_numpy( From 8f34b573be9066ea85aa950513e1fff06a98b3df Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 12:22:38 +0200 Subject: [PATCH 062/105] fix advection --- requirements-dev.txt | 1 + requirements.txt | 1 + tools/.pre-commit-config.yaml | 2 +- tools/pyproject.toml | 1 + tools/requirements-dev.txt | 1 + tools/requirements.txt | 1 + tools/src/icon4pytools/liskov/external/gt4py.py | 3 ++- tox.ini | 2 +- 8 files changed, 9 insertions(+), 3 deletions(-) diff --git a/requirements-dev.txt b/requirements-dev.txt index 85722a57b6..ae585c1cef 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,6 +3,7 @@ # Concretized dependencies in local repo -e ./common -e ./atm_dyn_iconam +-e ./advection # icon4pytools -e ./tools diff --git a/requirements.txt b/requirements.txt index d17dbd0ac3..22a84af692 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ # Concretized dependencies in local repo ./common ./atm_dyn_iconam +./advection # icon4pytools ./tools diff --git a/tools/.pre-commit-config.yaml b/tools/.pre-commit-config.yaml index 6c78812dfd..533cbae3f9 100644 --- a/tools/.pre-commit-config.yaml +++ b/tools/.pre-commit-config.yaml @@ -5,7 +5,7 @@ default_language_version: # Remove frozen version once we migrated away from tsa node: 17.9.1 minimum_pre_commit_version: 2.20.0 -exclude: "atm_dyn_iconam/.*" +exclude: "advection/.*", "atm_dyn_iconam/.*" repos: - repo: meta diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 24ff93a41c..b8a6e36535 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -27,6 +27,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ + "icon4py_advection", "icon4py_atm_dyn_iconam", "gt4py>=1.0.1", "icon4py_common", diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index 217c64d713..7b1ad57ba8 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,4 +1,5 @@ -r ../base-requirements-dev.txt -e ../common +-e ../advection -e ../atm_dyn_iconam -e . diff --git a/tools/requirements.txt b/tools/requirements.txt index b6eac2aa8a..e9c0e49500 100644 --- a/tools/requirements.txt +++ b/tools/requirements.txt @@ -1,4 +1,5 @@ -r ../base-requirements.txt ../common +../advection ../atm_dyn_iconam . diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 6996628825..755194875e 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -61,7 +61,8 @@ def _collect_icon4py_stencil(self, stencil_name: str) -> Program: err_counter += 1 if err_counter == len(self._STENCIL_PACKAGES): - raise UnknownStencilError(f"Did not find module: {stencil_name} in icon4pytools.") + error_msg = f"Did not find module: {stencil_name} in icon4pytools." + raise UnknownStencilError(error_msg) module_members = getmembers(module) found_stencil = [elt for elt in module_members if elt[0] == stencil_name] diff --git a/tox.ini b/tox.ini index 5c9ee6fd9d..1f63dc4c2c 100644 --- a/tox.ini +++ b/tox.ini @@ -16,7 +16,7 @@ setenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atm_dyn_iconam/src common/src pyutils/src testutils/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules advection/src atm_dyn_iconam/src common/src pyutils/src testutils/src pytest -v -s -n auto --cov --cov-append --ignore=tools/ commands_post = rm -Rf _reports/coverage_html From ff2885f48cf8346347a10cdb28753e715fdc6337 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 14:42:53 +0200 Subject: [PATCH 063/105] fix advection tests --- .../tests/test_face_val_ppm_stencil_01.py | 2 +- .../tests/test_face_val_ppm_stencil_02.py | 2 +- .../tests/test_face_val_ppm_stencil_05.py | 2 +- .../tests/test_hflx_limiter_mo_stencil_02.py | 2 +- .../tests/test_hflx_limiter_mo_stencil_03.py | 2 +- .../tests/test_hflx_limiter_mo_stencil_04.py | 2 +- .../tests/test_hflx_limiter_pd_stencil_01.py | 2 +- .../tests/test_hflx_limiter_pd_stencil_02.py | 2 +- advection/tests/test_hor_adv_stencil_01.py | 2 +- .../tests/test_rbf_intp_edge_stencil_01.py | 2 +- advection/tests/test_set_zero_c.py | 2 +- advection/tests/test_set_zero_c_k.py | 2 +- .../tests/test_step_advection_stencil_01.py | 2 +- .../tests/test_step_advection_stencil_02.py | 2 +- .../tests/test_step_advection_stencil_03.py | 2 +- .../tests/test_step_advection_stencil_04.py | 2 +- .../test_upwind_hflux_miura_stencil_01.py | 2 +- .../test_upwind_hflux_miura_stencil_02.py | 2 +- .../tests/test_upwind_vflux_ppm_stencil_01.py | 2 +- advection/tests/test_utils/__init__.py | 12 + advection/tests/test_utils/helpers.py | 106 +++++ advection/tests/test_utils/simple_mesh.py | 427 ++++++++++++++++++ advection/tests/test_vert_adv_stencil_01.py | 2 +- .../tests/test_vlimit_prbl_sm_stencil_01.py | 2 +- .../tests/test_vlimit_prbl_sm_stencil_02.py | 2 +- 25 files changed, 567 insertions(+), 22 deletions(-) create mode 100644 advection/tests/test_utils/__init__.py create mode 100644 advection/tests/test_utils/helpers.py create mode 100644 advection/tests/test_utils/simple_mesh.py diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 248dc991ac..37083960eb 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -18,7 +18,7 @@ from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import _shape, random_field, zero_field +from .test_utils.helpers import _shape, random_field, zero_field def face_val_ppm_stencil_01_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py index 56a82e1127..f3db55e678 100644 --- a/advection/tests/test_face_val_ppm_stencil_02.py +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -18,7 +18,7 @@ from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import _shape, random_field +from .test_utils.helpers import _shape, random_field def face_val_ppm_stencil_02_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/advection/tests/test_face_val_ppm_stencil_05.py index 96d17c85fc..546f135904 100644 --- a/advection/tests/test_face_val_ppm_stencil_05.py +++ b/advection/tests/test_face_val_ppm_stencil_05.py @@ -16,7 +16,7 @@ from icon4py.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def face_val_ppm_stencil_05_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py index 89f1e9ea78..8b85bc709f 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -19,7 +19,7 @@ ) from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import constant_field, random_field, zero_field +from .test_utils.helpers import constant_field, random_field, zero_field def hflx_limiter_mo_stencil_02_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py index c1b1c0b3e2..e6ee6d5065 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def hflx_limiter_mo_stencil_03_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/advection/tests/test_hflx_limiter_mo_stencil_04.py index 8c0436bcdd..abe11b37f9 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def hflx_limiter_mo_stencil_04_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/advection/tests/test_hflx_limiter_pd_stencil_01.py index 25311d1bc3..42d0e919c7 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -19,7 +19,7 @@ ) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import as_1D_sparse_field, random_field, zero_field +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field def hflx_limiter_pd_stencil_01_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index 8ef646fe81..2b2fa05efe 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import constant_field, random_field +from .test_utils.helpers import constant_field, random_field def hflx_limiter_pd_stencil_02_numpy( diff --git a/advection/tests/test_hor_adv_stencil_01.py b/advection/tests/test_hor_adv_stencil_01.py index 1274e7f2fd..2bf2b8fc70 100644 --- a/advection/tests/test_hor_adv_stencil_01.py +++ b/advection/tests/test_hor_adv_stencil_01.py @@ -17,7 +17,7 @@ from icon4py.advection.hor_adv_stencil_01 import hor_adv_stencil_01 from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import as_1D_sparse_field, random_field +from .test_utils.helpers import as_1D_sparse_field, random_field def hor_adv_stencil_01_numpy( diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/advection/tests/test_rbf_intp_edge_stencil_01.py index 80e0a5f3a7..cdfe853be7 100644 --- a/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -16,7 +16,7 @@ from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def rbf_intp_edge_stencil_01_numpy( diff --git a/advection/tests/test_set_zero_c.py b/advection/tests/test_set_zero_c.py index 4d887612f0..dab4dd81ea 100644 --- a/advection/tests/test_set_zero_c.py +++ b/advection/tests/test_set_zero_c.py @@ -16,7 +16,7 @@ from icon4py.advection.set_zero_c import set_zero_c from icon4py.common.dimension import CellDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def test_set_zero_cell_k(): diff --git a/advection/tests/test_set_zero_c_k.py b/advection/tests/test_set_zero_c_k.py index 9d23cf8b6d..79c202af14 100644 --- a/advection/tests/test_set_zero_c_k.py +++ b/advection/tests/test_set_zero_c_k.py @@ -16,7 +16,7 @@ from icon4py.advection.set_zero_c_k import set_zero_c_k from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def test_set_zero_c_k(): diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py index 8c8d6960a0..8e07eec025 100644 --- a/advection/tests/test_step_advection_stencil_01.py +++ b/advection/tests/test_step_advection_stencil_01.py @@ -16,7 +16,7 @@ from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def step_advection_stencil_01_numpy( diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py index e735bc8a4a..0c5d4e4bb1 100644 --- a/advection/tests/test_step_advection_stencil_02.py +++ b/advection/tests/test_step_advection_stencil_02.py @@ -16,7 +16,7 @@ from icon4py.advection.step_advection_stencil_02 import step_advection_stencil_02 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def step_advection_stencil_02_numpy( diff --git a/advection/tests/test_step_advection_stencil_03.py b/advection/tests/test_step_advection_stencil_03.py index 33997d8b11..b029bfb688 100644 --- a/advection/tests/test_step_advection_stencil_03.py +++ b/advection/tests/test_step_advection_stencil_03.py @@ -16,7 +16,7 @@ from icon4py.advection.step_advection_stencil_03 import step_advection_stencil_03 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field +from .test_utils.helpers import random_field def step_advection_stencil_03_numpy( diff --git a/advection/tests/test_step_advection_stencil_04.py b/advection/tests/test_step_advection_stencil_04.py index fd99e37083..d836d41284 100644 --- a/advection/tests/test_step_advection_stencil_04.py +++ b/advection/tests/test_step_advection_stencil_04.py @@ -16,7 +16,7 @@ from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def step_advection_stencil_04_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/advection/tests/test_upwind_hflux_miura_stencil_01.py index dbdd357c78..4bc67bcc54 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def upwind_hflux_miura_stencil_01_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py index 13ce1b5077..1724388fb5 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_02.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -19,7 +19,7 @@ ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import as_1D_sparse_field, random_field, zero_field +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field def upwind_hflux_miura_stencil_02_numpy( diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 96f2990e45..87be503483 100644 --- a/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def upwind_vflux_ppm_stencil_01_numpy( diff --git a/advection/tests/test_utils/__init__.py b/advection/tests/test_utils/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/advection/tests/test_utils/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/advection/tests/test_utils/helpers.py b/advection/tests/test_utils/helpers.py new file mode 100644 index 0000000000..f85bbf8f58 --- /dev/null +++ b/advection/tests/test_utils/helpers.py @@ -0,0 +1,106 @@ +# 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 typing import Optional + +import numpy as np +import numpy.typing as npt +from gt4py.next import common as gt_common +from gt4py.next.iterator import embedded as it_embedded + +from .simple_mesh import SimpleMesh + + +def _shape( + mesh, + *dims: gt_common.Dimension, + extend: Optional[dict[gt_common.Dimension, int]] = None, +): + if extend is None: + extend = {} + for d in dims: + if d not in extend.keys(): + extend[d] = 0 + return tuple(mesh.size[dim] + extend[dim] for dim in dims) + + +def random_mask( + mesh: SimpleMesh, + *dims: gt_common.Dimension, + dtype: Optional[npt.DTypeLike] = None, + extend: Optional[dict[gt_common.Dimension, int]] = None, +) -> it_embedded.MutableLocatedField: + shape = _shape(mesh, *dims, extend=extend) + arr = np.full(shape, False).flatten() + arr[: int(arr.size * 0.5)] = True + np.random.shuffle(arr) + arr = np.reshape(arr, newshape=shape) + if dtype: + arr = arr.astype(dtype) + return it_embedded.np_as_located_field(*dims)(arr) + + +def random_field( + mesh, + *dims, + low: float = -1.0, + high: float = 1.0, + extend: Optional[dict[gt_common.Dimension, int]] = None, +) -> it_embedded.MutableLocatedField: + return it_embedded.np_as_located_field(*dims)( + np.random.default_rng().uniform( + low=low, high=high, size=_shape(mesh, *dims, extend=extend) + ) + ) + + +def zero_field( + mesh: SimpleMesh, + *dims: gt_common.Dimension, + dtype=float, + extend: Optional[dict[gt_common.Dimension, int]] = None, +) -> it_embedded.MutableLocatedField: + return it_embedded.np_as_located_field(*dims)( + np.zeros(shape=_shape(mesh, *dims, extend=extend), dtype=dtype) + ) + + +def constant_field( + mesh: SimpleMesh, value: float, *dims: gt_common.Dimension, dtype=float +) -> it_embedded.MutableLocatedField: + return it_embedded.np_as_located_field(*dims)( + value * np.ones(shape=tuple(map(lambda x: mesh.size[x], dims)), dtype=dtype) + ) + + +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.""" + old_shape = np.asarray(field).shape + assert len(old_shape) == 2 + 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: + """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) diff --git a/advection/tests/test_utils/simple_mesh.py b/advection/tests/test_utils/simple_mesh.py new file mode 100644 index 0000000000..5b66aeed19 --- /dev/null +++ b/advection/tests/test_utils/simple_mesh.py @@ -0,0 +1,427 @@ +# 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 dataclasses import dataclass + +import numpy as np +from gt4py.next.iterator.embedded import NeighborTableOffsetProvider + +from icon4py.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CEDim, + CellDim, + E2C2EDim, + E2C2EODim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2CDim, + V2EDim, + VertexDim, +) + + +# periodic +# +# 0v---0e-- 1v---3e-- 2v---6e-- 0v +# | \ 0c | \ 1c | \2c +# | \1e | \4e | \7e +# |2e \ |5e \ |8e \ +# | 3c \ | 4c \ | 5c\ +# 3v---9e-- 4v--12e-- 5v--15e-- 3v +# | \ 6c | \ 7c | \ 8c +# | \10e | \13e | \16e +# |11e \ |14e \ |17e \ +# | 9c \ | 10c \ | 11c \ +# 6v--18e-- 7v--21e-- 8v--24e-- 6v +# | \12c | \ 13c | \ 14c +# | \19e | \22e | \25e +# |20e \ |23e \ |26e \ +# | 15c \ | 16c \ | 17c \ +# 0v 1v 2v 0v + + +@dataclass +class SimpleMeshData: + e2c2v_table = np.asarray( + [ + [0, 1, 4, 6], # 0 + [0, 4, 1, 3], # 1 + [0, 3, 4, 2], # 2 + [1, 2, 5, 7], # 3 + [1, 5, 2, 4], # 4 + [1, 4, 5, 0], # 5 + [2, 0, 3, 8], # 6 + [2, 3, 5, 0], # 7 + [2, 5, 1, 3], # 8 + [3, 4, 0, 7], # 9 + [3, 7, 4, 6], # 10 + [3, 6, 7, 5], # 11 + [4, 5, 8, 1], # 12 + [4, 8, 7, 5], # 13 + [4, 7, 3, 8], # 14 + [5, 3, 6, 2], # 15 + [6, 5, 3, 8], # 16 + [8, 5, 6, 4], # 17 + [6, 7, 3, 1], # 18 + [6, 1, 7, 0], # 19 + [6, 0, 1, 8], # 20 + [7, 8, 2, 4], # 21 + [7, 2, 8, 1], # 22 + [7, 1, 2, 6], # 23 + [8, 6, 0, 5], # 24 + [8, 0, 6, 2], # 25 + [8, 2, 0, 6], # 26 + ] + ) + + e2c_table = np.asarray( + [ + [0, 15], + [0, 3], + [3, 2], + [1, 16], + [1, 4], + [0, 4], + [2, 17], + [2, 5], + [1, 5], + [3, 6], + [6, 9], + [9, 8], + [4, 7], + [7, 10], + [6, 10], + [5, 8], + [8, 11], + [7, 11], + [9, 12], + [12, 15], + [15, 14], + [10, 13], + [13, 16], + [12, 16], + [11, 14], + [14, 17], + [13, 17], + ] + ) + + e2v_table = np.asarray( + [ + [0, 1], + [0, 4], + [0, 3], + [1, 2], + [1, 5], + [1, 4], + [2, 0], + [2, 3], + [2, 5], + [3, 4], + [3, 7], + [3, 6], + [4, 5], + [4, 8], + [4, 7], + [5, 3], + [5, 6], + [5, 8], + [6, 7], + [6, 1], + [6, 0], + [7, 8], + [7, 2], + [7, 1], + [8, 6], + [8, 0], + [8, 2], + ] + ) + + e2c2e_table = np.asarray( + [ + [1, 5, 19, 20], + [0, 5, 2, 9], + [1, 9, 6, 7], + [4, 8, 22, 23], + [3, 8, 5, 12], + [0, 1, 4, 12], + [7, 2, 25, 26], + [6, 2, 8, 15], + [3, 4, 7, 15], + [1, 2, 10, 14], + [9, 14, 11, 18], + [10, 18, 15, 16], + [4, 5, 13, 17], + [12, 17, 14, 21], + [9, 10, 13, 21], + [7, 8, 16, 11], + [15, 11, 17, 24], + [12, 13, 16, 24], + [10, 11, 19, 23], + [18, 23, 20, 0], + [19, 0, 24, 25], + [13, 14, 22, 26], + [21, 26, 23, 3], + [18, 19, 22, 3], + [16, 17, 25, 20], + [24, 20, 26, 6], + [25, 6, 21, 22], + ] + ) + + e2c2eO_table = np.asarray( + [ + [0, 1, 5, 19, 20], + [0, 1, 5, 2, 9], + [1, 2, 9, 6, 7], + [3, 4, 8, 22, 23], + [3, 4, 8, 5, 12], + [0, 1, 5, 4, 12], + [6, 7, 2, 25, 26], + [6, 7, 2, 8, 15], + [3, 4, 8, 7, 15], + [1, 2, 9, 10, 14], + [9, 10, 14, 11, 18], + [10, 11, 18, 15, 16], + [4, 5, 12, 13, 17], + [12, 13, 17, 14, 21], + [9, 10, 14, 13, 21], + [7, 8, 15, 16, 11], + [15, 16, 11, 17, 24], + [12, 13, 17, 16, 24], + [10, 11, 18, 19, 23], + [18, 19, 23, 20, 0], + [19, 20, 0, 24, 25], + [13, 14, 21, 22, 26], + [21, 22, 26, 23, 3], + [18, 19, 23, 22, 3], + [16, 17, 24, 25, 20], + [24, 25, 20, 26, 6], + [25, 26, 6, 21, 22], + ] + ) + + c2e_table = np.asarray( + [ + [0, 1, 5], # cell 0 + [3, 4, 8], # cell 1 + [6, 7, 2], # cell 2 + [1, 2, 9], # cell 3 + [4, 5, 12], # cell 4 + [7, 8, 15], # cell 5 + [9, 10, 14], # cell 6 + [12, 13, 17], # cell 7 + [15, 16, 11], # cell 8 + [10, 11, 18], # cell 9 + [13, 14, 21], # cell 10 + [16, 17, 24], # cell 11 + [18, 19, 23], # cell 12 + [21, 22, 26], # cell 13 + [24, 25, 20], # cell 14 + [19, 20, 0], # cell 15 + [22, 23, 3], # cell 16 + [25, 26, 6], # cell 17 + ] + ) + + v2c_table = np.asarray( + [ + [17, 14, 3, 0, 2, 15], + [0, 4, 1, 12, 16, 15], + [1, 5, 2, 16, 13, 17], + [3, 6, 9, 5, 8, 2], + [6, 10, 7, 4, 0, 3], + [7, 11, 8, 5, 1, 4], + [9, 12, 15, 8, 11, 14], + [12, 16, 13, 10, 6, 9], + [13, 17, 14, 11, 7, 10], + ] + ) + + v2e_table = np.asarray( + [ + [0, 1, 2, 6, 25, 20], + [3, 4, 5, 0, 23, 19], + [6, 7, 8, 3, 22, 26], + [9, 10, 11, 15, 7, 2], + [12, 13, 14, 9, 1, 5], + [15, 16, 17, 12, 4, 8], + [18, 19, 20, 24, 16, 11], + [21, 22, 23, 18, 10, 14], + [24, 25, 26, 21, 13, 17], + ] + ) + + diamond_table = np.asarray( + [ + [0, 1, 4, 6], # 0 + [0, 4, 1, 3], + [0, 3, 4, 2], + [1, 2, 5, 7], # 3 + [1, 5, 2, 4], + [1, 4, 5, 0], + [2, 0, 3, 8], # 6 + [2, 3, 0, 5], + [2, 5, 1, 3], + [3, 4, 0, 7], # 9 + [3, 7, 4, 6], + [3, 6, 5, 7], + [4, 5, 1, 8], # 12 + [4, 8, 5, 7], + [4, 7, 3, 8], + [5, 3, 2, 6], # 15 + [5, 6, 3, 8], + [5, 8, 4, 6], + [6, 7, 3, 1], # 18 + [6, 1, 7, 0], + [6, 0, 1, 8], + [7, 8, 4, 2], # 21 + [7, 2, 8, 1], + [7, 1, 6, 2], + [8, 6, 5, 0], # 24 + [8, 0, 6, 2], + [8, 2, 7, 0], + ] + ) + + c2e2cO_table = np.asarray( + [ + [15, 4, 3, 0], + [16, 5, 4, 1], + [17, 3, 5, 2], + [0, 6, 2, 3], + [1, 7, 0, 4], + [2, 8, 1, 5], + [3, 10, 9, 6], + [4, 11, 10, 7], + [5, 9, 11, 8], + [6, 12, 8, 9], + [7, 13, 6, 10], + [8, 14, 7, 11], + [9, 16, 15, 12], + [10, 17, 16, 13], + [11, 15, 17, 14], + [12, 0, 14, 15], + [13, 1, 12, 16], + [14, 2, 13, 17], + ] + ) + + c2e2c_table = np.asarray( + [ + [15, 4, 3], + [16, 5, 4], + [17, 3, 5], + [0, 6, 2], + [1, 7, 0], + [2, 8, 1], + [3, 10, 9], + [4, 11, 10], + [5, 9, 11], + [6, 12, 8], + [7, 13, 6], + [8, 14, 7], + [9, 16, 15], + [10, 17, 16], + [11, 15, 17], + [12, 0, 14], + [13, 1, 12], + [14, 2, 13], + ] + ) + + +class SimpleMesh: + _DEFAULT_K_LEVEL = 10 + + def __init__(self, k_level: int = _DEFAULT_K_LEVEL): + self.diamond_arr = SimpleMeshData.diamond_table + self.e2c = SimpleMeshData.e2c_table + self.e2v = SimpleMeshData.e2v_table + self.c2e = SimpleMeshData.c2e_table + self.c2e2cO = SimpleMeshData.c2e2cO_table + self.c2e2c = SimpleMeshData.c2e2c_table + self.e2c2eO = SimpleMeshData.e2c2eO_table + self.e2c2e = SimpleMeshData.e2c2e_table + self.e2c2v = SimpleMeshData.e2c2v_table + self.v2c = SimpleMeshData.v2c_table + self.v2e = SimpleMeshData.v2e_table + self.n_e2c = self.e2c.shape[1] + self.n_e2v = self.e2v.shape[1] + self.n_c2e = self.c2e.shape[1] + self.n_c2e2cO = self.c2e2cO.shape[1] + self.n_c2e2c = self.c2e2c.shape[1] + self.n_e2c2eO = self.e2c2eO.shape[1] + self.n_e2c2e = self.e2c2e.shape[1] + self.n_e2c2v = self.e2c2v.shape[1] + self.n_v2c = self.v2c.shape[1] + self.n_v2e = self.v2e.shape[1] + self.n_cells = self.c2e.shape[0] + self.n_edges = 27 + self.n_vertices = 9 + self.k_level = k_level + self.size = { + CellDim: self.n_cells, + EdgeDim: self.n_edges, + E2VDim: self.n_e2v, + E2CDim: self.n_e2c, + C2EDim: self.n_c2e, + C2E2CODim: self.n_c2e2cO, + C2E2CDim: self.n_c2e2c, + E2C2EODim: self.n_e2c2eO, + E2C2EDim: self.n_e2c2e, + V2CDim: self.n_v2c, + KDim: self.k_level, + VertexDim: self.n_vertices, + V2EDim: self.n_v2e, + CEDim: self.n_cells * self.n_c2e, + E2C2VDim: self.n_e2c2v, + ECVDim: self.n_edges * self.n_e2c2v, + } + + def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e, CellDim, EdgeDim, self.n_c2e) + + def get_c2e2cO_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e2cO, CellDim, CellDim, self.n_c2e2cO) + + def get_c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e2c, CellDim, CellDim, self.n_c2e2c) + + def get_e2c2eO_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.e2c2eO, EdgeDim, EdgeDim, self.n_e2c2eO) + + def get_e2c2e_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.e2c2e, EdgeDim, EdgeDim, self.n_e2c2e) + + def get_v2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.v2c, VertexDim, CellDim, self.n_v2c) + + def get_v2e_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.v2e, VertexDim, EdgeDim, self.n_v2e) + + def get_e2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.e2c, EdgeDim, CellDim, self.n_e2c) + + def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.e2v, EdgeDim, VertexDim, self.n_e2v) + + def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py index 15416c3d58..2782dea382 100644 --- a/advection/tests/test_vert_adv_stencil_01.py +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -16,7 +16,7 @@ from icon4py.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def vert_adv_stencil_01_numpy( diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 0a20d24798..c38a05b552 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field from gt4py.next.ffront.fbuiltins import int32 diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/advection/tests/test_vlimit_prbl_sm_stencil_02.py index cef8b8d589..263628aec0 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, KDim from .test_utils.simple_mesh import SimpleMesh -from .test_utils.utils import random_field, zero_field +from .test_utils.helpers import random_field, zero_field def v_limit_prbl_sm_stencil_01_numpy( From 5f5fd4062aca4d64eadc2bc7b581c9394cce6e3c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 14:50:24 +0200 Subject: [PATCH 064/105] fix style --- .../icon4py/advection/v_limit_prbl_sm_stencil_01.py | 6 +++--- .../icon4py/advection/v_limit_prbl_sm_stencil_02.py | 5 ++--- advection/tests/test_face_val_ppm_stencil_01.py | 3 ++- advection/tests/test_face_val_ppm_stencil_02.py | 3 ++- advection/tests/test_face_val_ppm_stencil_05.py | 3 ++- advection/tests/test_hflx_limiter_mo_stencil_02.py | 3 ++- advection/tests/test_hflx_limiter_mo_stencil_03.py | 3 ++- advection/tests/test_hflx_limiter_mo_stencil_04.py | 3 ++- advection/tests/test_hflx_limiter_pd_stencil_01.py | 3 ++- advection/tests/test_hflx_limiter_pd_stencil_02.py | 3 ++- advection/tests/test_hor_adv_stencil_01.py | 3 ++- advection/tests/test_rbf_intp_edge_stencil_01.py | 3 ++- advection/tests/test_set_zero_c.py | 3 ++- advection/tests/test_set_zero_c_k.py | 3 ++- advection/tests/test_step_advection_stencil_01.py | 3 ++- advection/tests/test_step_advection_stencil_02.py | 3 ++- advection/tests/test_step_advection_stencil_03.py | 3 ++- advection/tests/test_step_advection_stencil_04.py | 3 ++- advection/tests/test_upwind_hflux_miura_stencil_01.py | 3 ++- advection/tests/test_upwind_hflux_miura_stencil_02.py | 3 ++- advection/tests/test_upwind_vflux_ppm_stencil_01.py | 3 ++- advection/tests/test_vert_adv_stencil_01.py | 3 ++- advection/tests/test_vlimit_prbl_sm_stencil_01.py | 10 +++++----- advection/tests/test_vlimit_prbl_sm_stencil_02.py | 3 ++- atm_dyn_iconam/tests/test_compute_airmass.py | 3 ++- 25 files changed, 54 insertions(+), 33 deletions(-) diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py index d4101936a5..3bd9bf5c1b 100644 --- a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.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, FieldOffset, abs, minimum, where, int32 +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, abs, int32, where from icon4py.common.dimension import CellDim, KDim @@ -29,8 +29,8 @@ def _v_limit_prbl_sm_stencil_01( z_delta = p_face - p_face(Koff[1]) z_a6i = 6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) - l_limit = where(abs(z_delta) < -1.0*z_a6i, int32(1), int32(0)) - + l_limit = where(abs(z_delta) < -1.0 * z_a6i, int32(1), int32(0)) + return l_limit diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py index 30b9cdadea..7315b58d3d 100644 --- a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py +++ b/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.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, FieldOffset, abs, minimum, where, int32 +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, int32, minimum, where from icon4py.common.dimension import CellDim, KDim @@ -27,9 +27,8 @@ def _v_limit_prbl_sm_stencil_02( p_cc: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - q_face_up, q_face_low = where( - l_limit!=int32(0), + l_limit != int32(0), where( (p_cc < minimum(p_face, p_face(Koff[1]))), (p_cc, p_cc), diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/advection/tests/test_face_val_ppm_stencil_01.py index 37083960eb..610a3b00d0 100644 --- a/advection/tests/test_face_val_ppm_stencil_01.py +++ b/advection/tests/test_face_val_ppm_stencil_01.py @@ -17,8 +17,9 @@ from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import _shape, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_01_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/advection/tests/test_face_val_ppm_stencil_02.py index f3db55e678..7062097cc7 100644 --- a/advection/tests/test_face_val_ppm_stencil_02.py +++ b/advection/tests/test_face_val_ppm_stencil_02.py @@ -17,8 +17,9 @@ from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import _shape, random_field +from .test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_02_numpy( diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/advection/tests/test_face_val_ppm_stencil_05.py index 546f135904..2dc9e1ace0 100644 --- a/advection/tests/test_face_val_ppm_stencil_05.py +++ b/advection/tests/test_face_val_ppm_stencil_05.py @@ -15,8 +15,9 @@ from icon4py.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_05_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py index 8b85bc709f..62d2339138 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -18,8 +18,9 @@ hflx_limiter_mo_stencil_02, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import constant_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_02_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py index e6ee6d5065..b1935fca28 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -17,8 +17,9 @@ hflx_limiter_mo_stencil_03_min_max, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_03_numpy( diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/advection/tests/test_hflx_limiter_mo_stencil_04.py index abe11b37f9..20b0008aa6 100644 --- a/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -17,8 +17,9 @@ hflx_limiter_mo_stencil_04, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_04_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/advection/tests/test_hflx_limiter_pd_stencil_01.py index 42d0e919c7..df8bbfa291 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -18,8 +18,9 @@ hflx_limiter_pd_stencil_01, ) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def hflx_limiter_pd_stencil_01_numpy( diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index 2b2fa05efe..550133e0e8 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -17,8 +17,9 @@ hflx_limiter_pd_stencil_02, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import constant_field, random_field +from .test_utils.simple_mesh import SimpleMesh def hflx_limiter_pd_stencil_02_numpy( diff --git a/advection/tests/test_hor_adv_stencil_01.py b/advection/tests/test_hor_adv_stencil_01.py index 2bf2b8fc70..e009c9a04d 100644 --- a/advection/tests/test_hor_adv_stencil_01.py +++ b/advection/tests/test_hor_adv_stencil_01.py @@ -16,8 +16,9 @@ from icon4py.advection.hor_adv_stencil_01 import hor_adv_stencil_01 from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import as_1D_sparse_field, random_field +from .test_utils.simple_mesh import SimpleMesh def hor_adv_stencil_01_numpy( diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/advection/tests/test_rbf_intp_edge_stencil_01.py index cdfe853be7..f998afba67 100644 --- a/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -15,8 +15,9 @@ from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def rbf_intp_edge_stencil_01_numpy( diff --git a/advection/tests/test_set_zero_c.py b/advection/tests/test_set_zero_c.py index dab4dd81ea..89ae27a48e 100644 --- a/advection/tests/test_set_zero_c.py +++ b/advection/tests/test_set_zero_c.py @@ -15,8 +15,9 @@ from icon4py.advection.set_zero_c import set_zero_c from icon4py.common.dimension import CellDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def test_set_zero_cell_k(): diff --git a/advection/tests/test_set_zero_c_k.py b/advection/tests/test_set_zero_c_k.py index 79c202af14..60ab010379 100644 --- a/advection/tests/test_set_zero_c_k.py +++ b/advection/tests/test_set_zero_c_k.py @@ -15,8 +15,9 @@ from icon4py.advection.set_zero_c_k import set_zero_c_k from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def test_set_zero_c_k(): diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py index 8e07eec025..7cab0293ee 100644 --- a/advection/tests/test_step_advection_stencil_01.py +++ b/advection/tests/test_step_advection_stencil_01.py @@ -15,8 +15,9 @@ from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def step_advection_stencil_01_numpy( diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py index 0c5d4e4bb1..030d190255 100644 --- a/advection/tests/test_step_advection_stencil_02.py +++ b/advection/tests/test_step_advection_stencil_02.py @@ -15,8 +15,9 @@ from icon4py.advection.step_advection_stencil_02 import step_advection_stencil_02 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def step_advection_stencil_02_numpy( diff --git a/advection/tests/test_step_advection_stencil_03.py b/advection/tests/test_step_advection_stencil_03.py index b029bfb688..d3aac81733 100644 --- a/advection/tests/test_step_advection_stencil_03.py +++ b/advection/tests/test_step_advection_stencil_03.py @@ -15,8 +15,9 @@ from icon4py.advection.step_advection_stencil_03 import step_advection_stencil_03 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field +from .test_utils.simple_mesh import SimpleMesh def step_advection_stencil_03_numpy( diff --git a/advection/tests/test_step_advection_stencil_04.py b/advection/tests/test_step_advection_stencil_04.py index d836d41284..0b7178a13e 100644 --- a/advection/tests/test_step_advection_stencil_04.py +++ b/advection/tests/test_step_advection_stencil_04.py @@ -15,8 +15,9 @@ from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def step_advection_stencil_04_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/advection/tests/test_upwind_hflux_miura_stencil_01.py index 4bc67bcc54..868ffc38dc 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -17,8 +17,9 @@ upwind_hflux_miura_stencil_01, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_stencil_01_numpy( diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py index 1724388fb5..d748eae380 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_02.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -18,8 +18,9 @@ upwind_hflux_miura_stencil_02, ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_stencil_02_numpy( diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 87be503483..37edce8962 100644 --- a/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -17,8 +17,9 @@ upwind_vflux_ppm_stencil_01, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def upwind_vflux_ppm_stencil_01_numpy( diff --git a/advection/tests/test_vert_adv_stencil_01.py b/advection/tests/test_vert_adv_stencil_01.py index 2782dea382..9b6261745f 100644 --- a/advection/tests/test_vert_adv_stencil_01.py +++ b/advection/tests/test_vert_adv_stencil_01.py @@ -15,8 +15,9 @@ from icon4py.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def vert_adv_stencil_01_numpy( diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/advection/tests/test_vlimit_prbl_sm_stencil_01.py index c38a05b552..e0e475bc0c 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -12,15 +12,15 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 from icon4py.advection.v_limit_prbl_sm_stencil_01 import ( v_limit_prbl_sm_stencil_01, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh -from .test_utils.helpers import random_field, zero_field -from gt4py.next.ffront.fbuiltins import int32 +from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def v_limit_prbl_sm_stencil_01_numpy( @@ -31,7 +31,7 @@ def v_limit_prbl_sm_stencil_01_numpy( z_delta = p_face[:, :-1] - p_face[:, 1:] z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - l_limit = np.where( np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) return l_limit @@ -40,7 +40,7 @@ def test_v_limit_prbl_sm_stencil_01(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - l_limit = zero_field(mesh, CellDim, KDim, dtype=in32) + l_limit = zero_field(mesh, CellDim, KDim, dtype=int32) l_limit_ref = v_limit_prbl_sm_stencil_01_numpy( np.asarray(p_face), diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/advection/tests/test_vlimit_prbl_sm_stencil_02.py index 263628aec0..ae31b65768 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -17,8 +17,9 @@ v_limit_prbl_sm_stencil_01, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh def v_limit_prbl_sm_stencil_01_numpy( diff --git a/atm_dyn_iconam/tests/test_compute_airmass.py b/atm_dyn_iconam/tests/test_compute_airmass.py index 1742984fcd..e67739f53f 100644 --- a/atm_dyn_iconam/tests/test_compute_airmass.py +++ b/atm_dyn_iconam/tests/test_compute_airmass.py @@ -15,8 +15,9 @@ from icon4py.atm_dyn_iconam.compute_airmass import compute_airmass from icon4py.common.dimension import CellDim, KDim -from .test_utils.simple_mesh import SimpleMesh + from .test_utils.helpers import random_field +from .test_utils.simple_mesh import SimpleMesh def compute_airmass_numpy( From 06f6e8d6f59196ece6e0760a6b1a6a1338da5437 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 15:15:16 +0200 Subject: [PATCH 065/105] fix test --- .../tests/test_vlimit_prbl_sm_stencil_02.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/advection/tests/test_vlimit_prbl_sm_stencil_02.py index ae31b65768..427a900f31 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -13,25 +13,26 @@ import numpy as np -from icon4py.advection.v_limit_prbl_sm_stencil_01 import ( - v_limit_prbl_sm_stencil_01, +from gt4py.next.ffront.fbuiltins import int32 + + +from icon4py.advection.v_limit_prbl_sm_stencil_02 import ( + v_limit_prbl_sm_stencil_02, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field +from .test_utils.helpers import random_field, zero_field, random_mask from .test_utils.simple_mesh import SimpleMesh -def v_limit_prbl_sm_stencil_01_numpy( +def v_limit_prbl_sm_stencil_02_numpy( + l_limit: np.array, p_face: np.array, p_cc: np.array, ): - z_delta = p_face[:, :-1] - p_face[:, 1:] - z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - q_face_up, q_face_low = np.where( - np.abs(z_delta) < -1 * z_a6i, + l_limit != int32(0), np.where( (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), (p_cc, p_cc), @@ -47,19 +48,22 @@ def v_limit_prbl_sm_stencil_01_numpy( return q_face_up, q_face_low -def test_v_limit_prbl_sm_stencil_01(): +def test_v_limit_prbl_sm_stencil_02(): mesh = SimpleMesh() + l_limit = random_mask(mesh, CellDim, KDim, dtype=int32) p_cc = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) p_face_up = zero_field(mesh, CellDim, KDim) p_face_low = zero_field(mesh, CellDim, KDim) - p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_01_numpy( + p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_02_numpy( + np.asarray(l_limit), np.asarray(p_face), np.asarray(p_cc), ) - v_limit_prbl_sm_stencil_01( + v_limit_prbl_sm_stencil_02( + l_limit, p_face, p_cc, p_face_up, From b4252745667858561706f126581e96967dd857e4 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 15:18:45 +0200 Subject: [PATCH 066/105] fix style --- advection/tests/test_vlimit_prbl_sm_stencil_02.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/advection/tests/test_vlimit_prbl_sm_stencil_02.py index 427a900f31..f98cb6ee90 100644 --- a/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -12,16 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np - from gt4py.next.ffront.fbuiltins import int32 - from icon4py.advection.v_limit_prbl_sm_stencil_02 import ( v_limit_prbl_sm_stencil_02, ) from icon4py.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field, random_mask +from .test_utils.helpers import random_field, random_mask, zero_field from .test_utils.simple_mesh import SimpleMesh From b6b406ac085ef55a45b11c88ddeaed4f4ab06165 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 13 Jun 2023 16:12:43 +0200 Subject: [PATCH 067/105] add init --- advection/__init__.py | 12 ++++++++++++ advection/tests/__init__.py | 12 ++++++++++++ 2 files changed, 24 insertions(+) create mode 100644 advection/__init__.py create mode 100644 advection/tests/__init__.py diff --git a/advection/__init__.py b/advection/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/advection/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/advection/tests/__init__.py b/advection/tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/advection/tests/__init__.py @@ -0,0 +1,12 @@ +# 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 095df4443488da8dab5759b2fcb0b89e4ca30564 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 10:25:56 +0200 Subject: [PATCH 068/105] fix stencil --- advection/src/icon4py/advection/face_val_ppm_stencil_01.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py index 53945a844e..2e4a651ebc 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_01.py @@ -66,9 +66,9 @@ def _face_val_ppm_stencil_01( vert_idx = broadcast(vert_idx, (CellDim, KDim)) z_slope = where( - vert_idx < elev, - _face_val_ppm_stencil_01a(p_cc, p_cellhgt_mc_now), + vert_idx == elev, _face_val_ppm_stencil_01b(p_cc, p_cellhgt_mc_now), + _face_val_ppm_stencil_01a(p_cc, p_cellhgt_mc_now), ) return z_slope From fd5bcb03b7cdd039067c3698a412956c12e9615c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 14:44:11 +0200 Subject: [PATCH 069/105] add optional feature to allow mergecopy for more then two stencils --- .../liskov/codegen/integration/exceptions.py | 3 + .../liskov/codegen/integration/generate.py | 66 ++++++++++++------- 2 files changed, 46 insertions(+), 23 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/exceptions.py b/tools/src/icon4pytools/liskov/codegen/integration/exceptions.py index cd7ca68398..21447c268b 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/exceptions.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/exceptions.py @@ -14,3 +14,6 @@ class UndeclaredFieldError(Exception): pass + +class UnmergedCopyError(Exception): + pass diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 9705435be3..f7c0099114 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -46,6 +46,8 @@ from icon4pytools.liskov.codegen.shared.generate import CodeGenerator from icon4pytools.liskov.codegen.shared.types import GeneratedCode from icon4pytools.liskov.external.metadata import CodeMetadata +from icon4pytools.liskov.codegen.integration.exceptions import UnmergedCopyError + logger = setup_logger(__name__) @@ -114,30 +116,48 @@ def _generate_start_stencil(self) -> None: stencil = self.interface.StartStencil[i] logger.info(f"Generating START statement for {stencil.name}") - try: - next_stencil = self.interface.StartStencil[i + 1] - except IndexError: - pass - - if stencil.mergecopy and next_stencil.mergecopy: - stencil = StartStencilData( - startln=stencil.startln, - name=stencil.name + "_" + next_stencil.name, - fields=stencil.fields + next_stencil.fields, - bounds=stencil.bounds, - acc_present=stencil.acc_present, - mergecopy=stencil.mergecopy, - copies=stencil.copies, - ) - i += 2 + if stencil.mergecopy: + + merged_name=stencil.name + merged_fields=stencil.fields.copy() + j = 1 + try: + next_stencil = self.interface.StartStencil[i + j] + except IndexError: + pass + + while( next_stencil.mergecopy ): + merged_name += " " + next_stencil.name + merged_fields += next_stencil.fields + j += 1 + try: + next_stencil = self.interface.StartStencil[i + j] + except IndexError: + pass + + stencil = StartStencilData( + startln=stencil.startln, + name=merged_name, + fields=merged_fields, + bounds=stencil.bounds, + acc_present=stencil.acc_present, + mergecopy=stencil.mergecopy, + copies=stencil.copies, + ) + + self._generate( + StartStencilStatement, + StartStencilStatementGenerator, + stencil.startln, + stencil_data=stencil, + profile=self.profile, + ) + + if j==1: + error_msg = f"{merged_name} cannot be merged with other stencil." + raise UnmergedCopyError(error_msg) + i += j - self._generate( - StartStencilStatement, - StartStencilStatementGenerator, - stencil.startln, - stencil_data=stencil, - profile=self.profile, - ) else: self._generate( StartStencilStatement, From de73562c078d244472a86367dede1807cc97072c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 14:45:57 +0200 Subject: [PATCH 070/105] split stencil --- .../advection/face_val_ppm_stencil_02a.py | 43 +++++++++++++++++++ .../advection/face_val_ppm_stencil_02b.py | 37 ++++++++++++++++ .../advection/face_val_ppm_stencil_02c.py | 39 +++++++++++++++++ 3 files changed, 119 insertions(+) create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_02a.py create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_02b.py create mode 100644 advection/src/icon4py/advection/face_val_ppm_stencil_02c.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py new file mode 100644 index 0000000000..bf93190784 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py @@ -0,0 +1,43 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) + + return p_face + + +@program +def face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02a( + p_cc, + p_cellhgt_mc_now, + out=p_face, + ) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py new file mode 100644 index 0000000000..c2e463ab3d --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py @@ -0,0 +1,37 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc + return p_face + + +@program +def face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02b( + p_cc, + out=p_face, + ) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py new file mode 100644 index 0000000000..617dd77976 --- /dev/null +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py @@ -0,0 +1,39 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.common.dimension import CellDim, KDim, Koff + + + +@field_operator +def _face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc(Koff[-1]) + return p_face + + + +@program +def face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02c( + p_cc, + out=p_face, + ) From f44a92ee2ce406b753fa7d7661bc772f625b24f5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 14:48:34 +0200 Subject: [PATCH 071/105] fix style --- advection/src/icon4py/advection/face_val_ppm_stencil_02a.py | 2 +- advection/src/icon4py/advection/face_val_ppm_stencil_02b.py | 4 ++-- advection/src/icon4py/advection/face_val_ppm_stencil_02c.py | 4 +--- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py index bf93190784..4d03e88a88 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02a.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, broadcast, int32, where +from gt4py.next.ffront.fbuiltins import Field from icon4py.common.dimension import CellDim, KDim, Koff diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py index c2e463ab3d..d280b24ec9 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where +from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.common.dimension import CellDim, KDim @field_operator diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py b/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py index 617dd77976..c34ad0342b 100644 --- a/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py +++ b/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py @@ -12,12 +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, broadcast, int32, where +from gt4py.next.ffront.fbuiltins import Field from icon4py.common.dimension import CellDim, KDim, Koff - @field_operator def _face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], @@ -27,7 +26,6 @@ def _face_val_ppm_stencil_02c( return p_face - @program def face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], From 1785a89ededa91bb6081a42e42d310cc960b59a1 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 15:30:22 +0200 Subject: [PATCH 072/105] add utest --- .../tests/test_face_val_ppm_stencil_02a.py | 63 +++++++++++++++++++ .../tests/test_face_val_ppm_stencil_02b.py | 50 +++++++++++++++ .../tests/test_face_val_ppm_stencil_02c.py | 51 +++++++++++++++ 3 files changed, 164 insertions(+) create mode 100644 advection/tests/test_face_val_ppm_stencil_02a.py create mode 100644 advection/tests/test_face_val_ppm_stencil_02b.py create mode 100644 advection/tests/test_face_val_ppm_stencil_02c.py diff --git a/advection/tests/test_face_val_ppm_stencil_02a.py b/advection/tests/test_face_val_ppm_stencil_02a.py new file mode 100644 index 0000000000..a657626574 --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_02a.py @@ -0,0 +1,63 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a +from icon4py.common.dimension import CellDim, KDim + +from .test_utils.helpers import _shape, random_field +from .test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02a_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, +): + + p_face = p_cc.copy() + + p_face[:, 1:] = p_cc[:, 1:] * ( + 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) + ) + ( + p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) + ) * ( + (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + + p_cc[:, :-1] + ) + + return p_face + + +def test_face_val_ppm_stencil_02a(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + + ref = face_val_ppm_stencil_02a_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + ) + + face_val_ppm_stencil_02a( + p_cc, + p_cellhgt_mc_now, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:,1:], p_face[:,1:]) diff --git a/advection/tests/test_face_val_ppm_stencil_02b.py b/advection/tests/test_face_val_ppm_stencil_02b.py new file mode 100644 index 0000000000..ce4511bc31 --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_02b.py @@ -0,0 +1,50 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b +from icon4py.common.dimension import CellDim, KDim + +from .test_utils.helpers import _shape, random_field +from .test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02b_numpy( + p_cc: np.array, +): + + p_face = p_cc.copy() + + return p_face + + +def test_face_val_ppm_stencil_02b(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + + ref = face_val_ppm_stencil_02b_numpy( + np.asarray(p_cc), + ) + + face_val_ppm_stencil_02b( + p_cc, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref, p_face) diff --git a/advection/tests/test_face_val_ppm_stencil_02c.py b/advection/tests/test_face_val_ppm_stencil_02c.py new file mode 100644 index 0000000000..4bfae8dc27 --- /dev/null +++ b/advection/tests/test_face_val_ppm_stencil_02c.py @@ -0,0 +1,51 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c +from icon4py.common.dimension import CellDim, KDim + +from .test_utils.helpers import _shape, random_field +from .test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02c_numpy( + p_cc: np.array, +): + + p_face = p_cc.copy() + + p_face[:, 1:] = p_cc[:, :-1] + + return p_face + + +def test_face_val_ppm_stencil_02c(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_02c_numpy( + np.asarray(p_cc), + ) + + face_val_ppm_stencil_02c( + p_cc, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:,1:], p_face[:,1:]) From e760e9fc2a71cffc797988a849f09418eee26ff2 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Jun 2023 15:38:06 +0200 Subject: [PATCH 073/105] fix style --- advection/tests/test_face_val_ppm_stencil_02a.py | 7 ++----- advection/tests/test_face_val_ppm_stencil_02b.py | 5 +---- advection/tests/test_face_val_ppm_stencil_02c.py | 6 ++---- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/advection/tests/test_face_val_ppm_stencil_02a.py b/advection/tests/test_face_val_ppm_stencil_02a.py index a657626574..a2735b03cc 100644 --- a/advection/tests/test_face_val_ppm_stencil_02a.py +++ b/advection/tests/test_face_val_ppm_stencil_02a.py @@ -12,13 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a from icon4py.common.dimension import CellDim, KDim -from .test_utils.helpers import _shape, random_field +from .test_utils.helpers import random_field from .test_utils.simple_mesh import SimpleMesh @@ -47,7 +45,6 @@ def test_face_val_ppm_stencil_02a(): p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim) - ref = face_val_ppm_stencil_02a_numpy( np.asarray(p_cc), np.asarray(p_cellhgt_mc_now), @@ -60,4 +57,4 @@ def test_face_val_ppm_stencil_02a(): offset_provider={"Koff": KDim}, ) - assert np.allclose(ref[:,1:], p_face[:,1:]) + assert np.allclose(ref[:, 1:], p_face[:, 1:]) diff --git a/advection/tests/test_face_val_ppm_stencil_02b.py b/advection/tests/test_face_val_ppm_stencil_02b.py index ce4511bc31..d8f5e3c1ee 100644 --- a/advection/tests/test_face_val_ppm_stencil_02b.py +++ b/advection/tests/test_face_val_ppm_stencil_02b.py @@ -12,13 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b from icon4py.common.dimension import CellDim, KDim -from .test_utils.helpers import _shape, random_field +from .test_utils.helpers import random_field from .test_utils.simple_mesh import SimpleMesh @@ -36,7 +34,6 @@ def test_face_val_ppm_stencil_02b(): p_cc = random_field(mesh, CellDim, KDim) p_face = random_field(mesh, CellDim, KDim) - ref = face_val_ppm_stencil_02b_numpy( np.asarray(p_cc), ) diff --git a/advection/tests/test_face_val_ppm_stencil_02c.py b/advection/tests/test_face_val_ppm_stencil_02c.py index 4bfae8dc27..9cc38cc5f5 100644 --- a/advection/tests/test_face_val_ppm_stencil_02c.py +++ b/advection/tests/test_face_val_ppm_stencil_02c.py @@ -12,13 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c from icon4py.common.dimension import CellDim, KDim -from .test_utils.helpers import _shape, random_field +from .test_utils.helpers import random_field from .test_utils.simple_mesh import SimpleMesh @@ -48,4 +46,4 @@ def test_face_val_ppm_stencil_02c(): offset_provider={"Koff": KDim}, ) - assert np.allclose(ref[:,1:], p_face[:,1:]) + assert np.allclose(ref[:, 1:], p_face[:, 1:]) From f327de28426c87f12140884b8693665ebde28e5d Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 28 Jun 2023 14:40:42 +0200 Subject: [PATCH 074/105] advection cycling routine addition + relative idx in btraj o1 modification --- .../upwind_hflux_miura_cycl_stencil_01.py | 59 +++++++++++++++ .../upwind_hflux_miura_cycl_stencil_02.py | 72 +++++++++++++++++++ .../upwind_hflux_miura_cycl_stencil_03a.py | 40 +++++++++++ .../upwind_hflux_miura_cycl_stencil_03b.py | 48 +++++++++++++ .../mo_advection_traj_btraj_compute_o1_dsl.py | 12 ++-- 5 files changed, 227 insertions(+), 4 deletions(-) create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py new file mode 100644 index 0000000000..5466fb2a03 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -0,0 +1,59 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_3_dsl: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + z_tracer_mflx_dsl = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl(E2C[1]), z_lsq_coeff_1_dsl(E2C[0])) + + distv_bary_1 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_dsl(E2C[1]), z_lsq_coeff_2_dsl(E2C[0])) + + distv_bary_2 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_dsl(E2C[1]), z_lsq_coeff_3_dsl(E2C[0])) + ) * p_mass_flx_e + + return z_tracer_mflx_dsl + +@program +def upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_3_dsl: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + z_tracer_mflx_dsl: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + out=(z_tracer_mflx_dsl) + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py new file mode 100644 index 0000000000..55fd2fb9b5 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -0,0 +1,72 @@ +# 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.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, where, neighbor_sum, int32, broadcast + +from icon4py.common.dimension import CellDim, KDim, C2EDim, EdgeDim, C2E + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_02( + nsub: int32, + p_mass_flx_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + z_rhofluxdiv_c: Field[[CellDim, KDim], float], + z_tracer_mflx: Field[[EdgeDim, KDim], float], + z_rho_now: Field[[CellDim, KDim], float], + z_tracer_now: Field[[CellDim, KDim], float], + z_dtsub: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + z_rhofluxdiv_c_out = neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) if nsub == int32(1) else z_rhofluxdiv_c + + z_fluxdiv_c_dsl = neighbor_sum(z_tracer_mflx(C2E) * geofac_div, axis=C2EDim) + + z_rho_new_dsl = ( z_rho_now + - z_dtsub * z_rhofluxdiv_c_out ) + + z_tracer_new_dsl = (( z_tracer_now + * z_rho_now + - z_dtsub * z_fluxdiv_c_dsl ) + / z_rho_new_dsl) + + return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + + +@program +def upwind_hflux_miura_cycl_stencil_02( + nsub: int32, + p_mass_flx_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + z_rhofluxdiv_c: Field[[CellDim, KDim], float], + z_tracer_mflx: Field[[EdgeDim, KDim], float], + z_rho_now: Field[[CellDim, KDim], float], + z_tracer_now: Field[[CellDim, KDim], float], + z_dtsub: float, + z_rhofluxdiv_c_out: Field[[CellDim, KDim], float], + z_fluxdiv_c_dsl: Field[[CellDim, KDim], float], + z_rho_new_dsl: Field[[CellDim, KDim], float], + z_tracer_new_dsl: Field[[CellDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_02( + nsub, + p_mass_flx_e, + geofac_div, + z_rhofluxdiv_c, + z_tracer_mflx, + z_rho_now, + z_tracer_now, + z_dtsub, + out=(z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl), + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py new file mode 100644 index 0000000000..71bbae9218 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -0,0 +1,40 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, float64 + +from icon4py.common.dimension import EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl)/float64(2) + return p_out_e + + +@program +def upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + out=(p_out_e) + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py new file mode 100644 index 0000000000..c654b484bd --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -0,0 +1,48 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.common.dimension import EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_03b( + p_ncycl: int32, + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_out_e = (z_tracer_mflx_1_dsl + + z_tracer_mflx_2_dsl + + z_tracer_mflx_3_dsl)/float(3) + return p_out_e + + +@program +def upwind_hflux_miura_cycl_stencil_03b( + p_ncycl: int32, + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_03b( + p_ncycl, + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + z_tracer_mflx_3_dsl, + out=(p_out_e) + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py index eb03d5ffaf..69ed968a7e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py @@ -38,7 +38,11 @@ def _mo_advection_traj_btraj_compute_o1_dsl( ]: lvn_pos = where(p_vn > 0.0, True, False) - p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) +# p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + + # Replace the absolute index of the upwind cell by index 0 and 1 + # Corresponding to neighboring cell 0 or 1 of an edge. + p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( @@ -67,7 +71,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 + return p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 @program @@ -82,7 +86,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_1: Field[[ECDim], float], primal_normal_cell_2: Field[[ECDim], float], dual_normal_cell_2: Field[[ECDim], float], - p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], p_distv_bary_1: Field[[EdgeDim, KDim], float], p_distv_bary_2: Field[[EdgeDim, KDim], float], @@ -100,5 +104,5 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2), + out=(p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2), ) From 3bc12daf5ea3147afd0fec09861e6477fc2484b7 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 08:36:04 +0200 Subject: [PATCH 075/105] adapt stencil btraj_compute_01 --- .../mo_advection_traj_btraj_compute_o1_dsl.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py index 69ed968a7e..b7161dc359 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py @@ -31,6 +31,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_2: Field[[ECDim], float], p_dthalf: float, ) -> tuple[ + Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], float], @@ -38,10 +39,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( ]: lvn_pos = where(p_vn > 0.0, True, False) -# p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) - - # Replace the absolute index of the upwind cell by index 0 and 1 - # Corresponding to neighboring cell 0 or 1 of an edge. + p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) @@ -71,7 +69,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 + return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 @program @@ -86,6 +84,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_1: Field[[ECDim], float], primal_normal_cell_2: Field[[ECDim], float], dual_normal_cell_2: Field[[ECDim], float], + p_cell_idx: Field[[EdgeDim, KDim], int32], p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], p_distv_bary_1: Field[[EdgeDim, KDim], float], @@ -104,5 +103,9 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=(p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2), + out=(p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_distv_bary_1, + p_distv_bary_2), ) From 0ec6e1eb8d30e8d0e9857d077b5f5d295a684396 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 09:16:26 +0200 Subject: [PATCH 076/105] adapt stencil upwind_hflux_miura_stencil_01.py --- .../upwind_hflux_miura_stencil_01.py | 53 ++++++++----------- 1 file changed, 21 insertions(+), 32 deletions(-) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py index 83fa538f14..e73f49649c 100644 --- a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py @@ -11,60 +11,49 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, where +from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C @field_operator def _upwind_hflux_miura_stencil_01( - p_vn: Field[[EdgeDim, KDim], float], - p_cc: Field[[CellDim, KDim], float], + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], distv_bary_1: Field[[EdgeDim, KDim], float], distv_bary_2: Field[[EdgeDim, KDim], float], - z_grad_1: Field[[CellDim, KDim], float], - z_grad_2: Field[[CellDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = where( - p_vn > 0.0, - ( - p_cc(E2C[0]) - + distv_bary_1 * z_grad_1(E2C[0]) - + distv_bary_2 * z_grad_2(E2C[0]) - ) - * p_mass_flx_e, - ( - p_cc(E2C[1]) - + distv_bary_1 * z_grad_1(E2C[1]) - + distv_bary_2 * z_grad_2(E2C[1]) - ) - * p_mass_flx_e, - ) - - return p_out_e + p_out_e = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) + + distv_bary_1 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) + + distv_bary_2 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) + ) * p_mass_flx_e + return p_out_e @program def upwind_hflux_miura_stencil_01( - p_vn: Field[[EdgeDim, KDim], float], - p_cc: Field[[CellDim, KDim], float], + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], distv_bary_1: Field[[EdgeDim, KDim], float], distv_bary_2: Field[[EdgeDim, KDim], float], - z_grad_1: Field[[CellDim, KDim], float], - z_grad_2: Field[[CellDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_out_e: Field[[EdgeDim, KDim], float], ): _upwind_hflux_miura_stencil_01( - p_vn, - p_cc, + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, distv_bary_1, distv_bary_2, - z_grad_1, - z_grad_2, p_mass_flx_e, - out=p_out_e, + cell_rel_idx_dsl, + out=(p_out_e) ) From 21628a3b661a85560aa788ae786110f2d5907a8f Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 14:04:09 +0200 Subject: [PATCH 077/105] add hybrid_stencil_01a and liskofy hybrid_stencil_02 --- .../hflux_ffsl_hybrid_stencil_01a.py | 109 ++++++++++++++++++ .../advection/hflux_ffsl_hybrid_stencil_02.py | 12 +- 2 files changed, 115 insertions(+), 6 deletions(-) create mode 100644 advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py new file mode 100644 index 0000000000..342b69e656 --- /dev/null +++ b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -0,0 +1,109 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C + + +@field_operator +def _hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum0_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float], + patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e_hybrid_1a = ( + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) * z_quad_vector_sum0_1 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) * z_quad_vector_sum0_2 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) * z_quad_vector_sum0_3 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_4(E2C[1]), z_lsq_coeff_4(E2C[0])) * z_quad_vector_sum0_4 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_5(E2C[1]), z_lsq_coeff_5(E2C[0])) * z_quad_vector_sum0_5 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_6(E2C[1]), z_lsq_coeff_6(E2C[0])) * z_quad_vector_sum0_6 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_7(E2C[1]), z_lsq_coeff_7(E2C[0])) * z_quad_vector_sum0_7 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_8(E2C[1]), z_lsq_coeff_8(E2C[0])) * z_quad_vector_sum0_8 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_9(E2C[1]), z_lsq_coeff_9(E2C[0])) * z_quad_vector_sum0_9 + + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0])) * z_quad_vector_sum0_10 + ) + + return p_out_e_hybrid_1a + +@program +def hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum0_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float], + patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_out_e_hybrid_1a: Field[[EdgeDim, KDim], float], +): + _hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum0_1, + z_quad_vector_sum0_2, + z_quad_vector_sum0_3, + z_quad_vector_sum0_4, + z_quad_vector_sum0_5, + z_quad_vector_sum0_6, + z_quad_vector_sum0_7, + z_quad_vector_sum0_8, + z_quad_vector_sum0_9, + z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl, + out=(p_out_e_hybrid_1a) + ) diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py index 70684d4baf..474ea49602 100644 --- a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py +++ b/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py @@ -19,25 +19,25 @@ @field_operator def _hflux_ffsl_hybrid_stencil_02( - p_out_e: Field[[EdgeDim, KDim], float], + p_out_e_hybrid_2: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], z_dreg_area: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = p_mass_flx_e * p_out_e / z_dreg_area + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area - return p_out_e + return p_out_e_hybrid_2 @program def hflux_ffsl_hybrid_stencil_02( - p_out_e: Field[[EdgeDim, KDim], float], + p_out_e_hybrid_2: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], z_dreg_area: Field[[EdgeDim, KDim], float], ): _hflux_ffsl_hybrid_stencil_02( - p_out_e, + p_out_e_hybrid_2, p_mass_flx_e, z_dreg_area, - out=p_out_e, + out=p_out_e_hybrid_2, ) From b244b55ab16608dc6a1fdf76a9fa8b2c21e10ae4 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 14:21:06 +0200 Subject: [PATCH 078/105] Adapt btraj_dreg_03 and modify short length condition --- .../advection/btraj_dreg_stencil_02.py | 10 +-- .../advection/btraj_dreg_stencil_03.py | 73 ++++++++++--------- 2 files changed, 44 insertions(+), 39 deletions(-) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py index 0b2a2f3630..1c4d116503 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py @@ -28,11 +28,11 @@ def _btraj_dreg_stencil_02( lvn_pos = where(p_vn >= 0.0, True, False) traj_length = sqrt(p_vn**2 + p_vt**2) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) - famask = where( - traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) + opt_famask_dsl = where( + traj_length > 0.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) ) - return famask + return opt_famask_dsl @program @@ -40,7 +40,7 @@ def btraj_dreg_stencil_02( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], edge_cell_length: Field[[ECDim], float], - famask: Field[[EdgeDim, KDim], int32], + opt_famask_dsl: Field[[EdgeDim, KDim], int32], p_dt: float, ): - _btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=famask) + _btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=opt_famask_dsl) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py index a43d1670f4..30283e3df3 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py @@ -5,11 +5,11 @@ # # 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 +# Free Software Foundation, either version 3 of the License, or any lat_dsler # 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 +# SPDX-License-Identifier: GPL-3.0-or-lat_dsler from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where @@ -38,6 +38,7 @@ def _btraj_dreg_stencil_03( lvn_sys_pos: Field[[EdgeDim, KDim], bool], p_dt: float, ) -> tuple[ + Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], float], @@ -55,6 +56,7 @@ def _btraj_dreg_stencil_03( lvn_pos = where(p_vn >= 0.0, True, False) p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) depart_pts_1_x = edge_verts_1_x - p_vn * p_dt @@ -91,42 +93,43 @@ def _btraj_dreg_stencil_03( dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) - p_coords_dreg_v_1_lon = ( + p_coords_dreg_v_1_lon_dsl = ( pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 ) - p_coords_dreg_v_2_lon = ( + p_coords_dreg_v_2_lon_dsl = ( pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 ) - p_coords_dreg_v_3_lon = ( + p_coords_dreg_v_3_lon_dsl = ( pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 ) - p_coords_dreg_v_4_lon = ( + p_coords_dreg_v_4_lon_dsl = ( pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 ) - p_coords_dreg_v_1_lat = ( + p_coords_dreg_v_1_lat_dsl = ( pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 ) - p_coords_dreg_v_2_lat = ( + p_coords_dreg_v_2_lat_dsl = ( pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 ) - p_coords_dreg_v_3_lat = ( + p_coords_dreg_v_3_lat_dsl = ( pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 ) - p_coords_dreg_v_4_lat = ( + p_coords_dreg_v_4_lat_dsl = ( pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 ) return ( p_cell_idx, + p_cell_rel_idx_dsl, p_cell_blk, - p_coords_dreg_v_1_lon, - p_coords_dreg_v_2_lon, - p_coords_dreg_v_3_lon, - p_coords_dreg_v_4_lon, - p_coords_dreg_v_1_lat, - p_coords_dreg_v_2_lat, - p_coords_dreg_v_3_lat, - p_coords_dreg_v_4_lat, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, ) @@ -151,15 +154,16 @@ def btraj_dreg_stencil_03( lvn_sys_pos: Field[[EdgeDim, KDim], bool], p_dt: float, p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], - p_coords_dreg_v_1_lon: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_2_lon: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_3_lon: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_4_lon: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_1_lat: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_2_lat: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_3_lat: Field[[EdgeDim, KDim], float], - p_coords_dreg_v_4_lat: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lat_dsl: Field[[EdgeDim, KDim], float], ): _btraj_dreg_stencil_03( p_vn, @@ -182,14 +186,15 @@ def btraj_dreg_stencil_03( p_dt, out=( p_cell_idx, + p_cell_rel_idx_dsl, p_cell_blk, - p_coords_dreg_v_1_lon, - p_coords_dreg_v_2_lon, - p_coords_dreg_v_3_lon, - p_coords_dreg_v_4_lon, - p_coords_dreg_v_1_lat, - p_coords_dreg_v_2_lat, - p_coords_dreg_v_3_lat, - p_coords_dreg_v_4_lat, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, ), ) From d035ebc747a7402e4d9d435df3847d3837321ba9 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 15:36:24 +0200 Subject: [PATCH 079/105] liskofy recon_lsq_cell_l/c_svd --- .../advection/recon_lsq_cell_c_svd_stencil.py | 98 +++++++++---------- .../advection/recon_lsq_cell_l_svd_stencil.py | 55 +++++++++++ 2 files changed, 104 insertions(+), 49 deletions(-) create mode 100644 advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py index 870ff50b81..03e5538bc0 100644 --- a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py +++ b/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py @@ -51,7 +51,7 @@ def _recon_lsq_cell_c_svd_stencil( Field[[CellDim, KDim], float], ]: - p_coeff_10 = ( + p_coeff_10_dsl = ( lsq_pseudoinv_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_9(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -62,7 +62,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_9(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_9(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_9 = ( + p_coeff_9_dsl = ( lsq_pseudoinv_8(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_8(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_8(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -73,7 +73,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_8(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_8(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_8 = ( + p_coeff_8_dsl = ( lsq_pseudoinv_7(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_7(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_7(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -84,7 +84,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_7(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_7(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_7 = ( + p_coeff_7_dsl = ( lsq_pseudoinv_6(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_6(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_6(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -95,7 +95,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_6(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_6(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_6 = ( + p_coeff_6_dsl = ( lsq_pseudoinv_5(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_5(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_5(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -106,7 +106,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_5(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_5(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_5 = ( + p_coeff_5_dsl = ( lsq_pseudoinv_4(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_4(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_4(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -117,7 +117,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_4(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_4(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_4 = ( + p_coeff_4_dsl = ( lsq_pseudoinv_3(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_3(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_3(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -128,7 +128,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_3(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_3(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_3 = ( + p_coeff_3_dsl = ( lsq_pseudoinv_2(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_2(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_2(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -139,7 +139,7 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_2(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + lsq_pseudoinv_2(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_2 = ( + p_coeff_2_dsl = ( lsq_pseudoinv_1(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_1(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + lsq_pseudoinv_1(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) @@ -151,29 +151,29 @@ def _recon_lsq_cell_c_svd_stencil( + lsq_pseudoinv_1(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) ) - p_coeff_1 = p_cc - ( - p_coeff_2 * lsq_moments_1 - + p_coeff_3 * lsq_moments_2 - + p_coeff_4 * lsq_moments_3 - + p_coeff_5 * lsq_moments_4 - + p_coeff_6 * lsq_moments_5 - + p_coeff_7 * lsq_moments_6 - + p_coeff_8 * lsq_moments_7 - + p_coeff_9 * lsq_moments_8 - + p_coeff_10 * lsq_moments_9 + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 ) return ( - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, ) @@ -198,16 +198,16 @@ def recon_lsq_cell_c_svd_stencil( lsq_moments_7: Field[[CellDim], float], lsq_moments_8: Field[[CellDim], float], lsq_moments_9: Field[[CellDim], float], - p_coeff_1: Field[[CellDim, KDim], float], - p_coeff_2: Field[[CellDim, KDim], float], - p_coeff_3: Field[[CellDim, KDim], float], - p_coeff_4: Field[[CellDim, KDim], float], - p_coeff_5: Field[[CellDim, KDim], float], - p_coeff_6: Field[[CellDim, KDim], float], - p_coeff_7: Field[[CellDim, KDim], float], - p_coeff_8: Field[[CellDim, KDim], float], - p_coeff_9: Field[[CellDim, KDim], float], - p_coeff_10: Field[[CellDim, KDim], float], + p_coeff_1_dsl: Field[[CellDim, KDim], float], + p_coeff_2_dsl: Field[[CellDim, KDim], float], + p_coeff_3_dsl: Field[[CellDim, KDim], float], + p_coeff_4_dsl: Field[[CellDim, KDim], float], + p_coeff_5_dsl: Field[[CellDim, KDim], float], + p_coeff_6_dsl: Field[[CellDim, KDim], float], + p_coeff_7_dsl: Field[[CellDim, KDim], float], + p_coeff_8_dsl: Field[[CellDim, KDim], float], + p_coeff_9_dsl: Field[[CellDim, KDim], float], + p_coeff_10_dsl: Field[[CellDim, KDim], float], ): _recon_lsq_cell_c_svd_stencil( p_cc, @@ -230,15 +230,15 @@ def recon_lsq_cell_c_svd_stencil( lsq_moments_8, lsq_moments_9, out=( - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, ), ) diff --git a/advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py b/advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py new file mode 100644 index 0000000000..2e574afd0c --- /dev/null +++ b/advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py @@ -0,0 +1,55 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim + + +@field_operator +def _recon_lsq_cell_l_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + p_coeff_1_dsl = p_cc + p_coeff_2_dsl = ( + lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + p_coeff_3_dsl = ( + lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + return p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl + + +@program +def recon_lsq_cell_l_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], + p_coeff_1_dsl: Field[[CellDim, KDim], float], + p_coeff_2_dsl: Field[[CellDim, KDim], float], + p_coeff_3_dsl: Field[[CellDim, KDim], float], +): + _recon_lsq_cell_l_svd_stencil( + p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl) + ) From 23a1bf234b51911854fdfd95c04c2f8e9c15fabc Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 20 Jul 2023 16:59:51 +0200 Subject: [PATCH 080/105] liskofy divide_flux_area_list_stencil_01/02.py --- .../divide_flux_area_list_stencil_01.py | 584 +++++++++--------- .../divide_flux_area_list_stencil_02.py | 224 +++---- 2 files changed, 404 insertions(+), 404 deletions(-) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py index 98a5ecd63e..d08b0d0bfa 100644 --- a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py @@ -127,14 +127,14 @@ def _divide_flux_area_list_stencil_01( ptr_v3_lon: Field[[ECDim], float], ptr_v3_lat: Field[[ECDim], float], tangent_orientation_dsl: Field[[EdgeDim], float], - dreg_patch0_1_x: Field[[EdgeDim, KDim], float], - dreg_patch0_1_y: Field[[EdgeDim, KDim], float], - dreg_patch0_2_x: Field[[EdgeDim, KDim], float], - dreg_patch0_2_y: Field[[EdgeDim, KDim], float], - dreg_patch0_3_x: Field[[EdgeDim, KDim], float], - dreg_patch0_3_y: Field[[EdgeDim, KDim], float], - dreg_patch0_4_x: Field[[EdgeDim, KDim], float], - dreg_patch0_4_y: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lat_dsl: Field[[EdgeDim, KDim], float], ) -> tuple[ Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], @@ -162,26 +162,26 @@ def _divide_flux_area_list_stencil_01( Field[[EdgeDim, KDim], float], ]: - arrival_pts_1_x = dreg_patch0_1_x - arrival_pts_1_y = dreg_patch0_1_y - arrival_pts_2_x = dreg_patch0_2_x - arrival_pts_2_y = dreg_patch0_2_y - depart_pts_1_x = dreg_patch0_4_x # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) - depart_pts_1_y = dreg_patch0_4_y - depart_pts_2_x = dreg_patch0_3_x - depart_pts_2_y = dreg_patch0_3_y + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl lvn_pos = where(p_vn >= 0.0, True, False) # get flux area departure-line segment - fl_line_p1_lon = depart_pts_1_x - fl_line_p1_lat = depart_pts_1_y - fl_line_p2_lon = depart_pts_2_x - fl_line_p2_lat = depart_pts_2_y + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl # get triangle edge 1 (A1V3) - tri_line1_p1_lon = arrival_pts_1_x - tri_line1_p1_lat = arrival_pts_1_y + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl tri_line1_p2_lon = where( lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), @@ -194,8 +194,8 @@ def _divide_flux_area_list_stencil_01( ) # get triangle edge 2 (A2V3) - tri_line2_p1_lon = arrival_pts_2_x - tri_line2_p1_lat = arrival_pts_2_y + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl tri_line2_p2_lon = where( lvn_pos, broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), @@ -258,150 +258,150 @@ def _divide_flux_area_list_stencil_01( ) # Case 1 - patch 0 - dreg_patch0_1_x = where(mask_case1, arrival_pts_1_x, dreg_patch0_1_x) - dreg_patch0_1_y = where(mask_case1, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where( - mask_case1, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x + dreg_patch0_1_lon_dsl = where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case1, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl ) - dreg_patch0_2_y = where( - mask_case1, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y + dreg_patch0_2_lat_dsl = where( + mask_case1, where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl ) - dreg_patch0_3_x = where(mask_case1, ps2_x, dreg_patch0_3_x) - dreg_patch0_3_y = where(mask_case1, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where( - mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x + dreg_patch0_3_lon_dsl = where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl ) - dreg_patch0_4_y = where( - mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y + dreg_patch0_4_lat_dsl = where( + mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl ) # Case 1 - patch 1 - dreg_patch1_1_x = where(mask_case1, arrival_pts_1_x, 0.0) - dreg_patch1_1_y = where(mask_case1, arrival_pts_1_y, 0.0) - dreg_patch1_4_x = where(mask_case1, arrival_pts_1_x, 0.0) - dreg_patch1_4_y = where(mask_case1, arrival_pts_1_y, 0.0) - dreg_patch1_2_x = where(mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_x), 0.0) - dreg_patch1_2_y = where(mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_y), 0.0) - dreg_patch1_3_x = where(mask_case1, where(lvn_sys_pos, depart_pts_1_x, ps1_x), 0.0) - dreg_patch1_3_y = where(mask_case1, where(lvn_sys_pos, depart_pts_1_y, ps1_y), 0.0) + dreg_patch1_1_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = where(mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0) + dreg_patch1_2_lat_vmask = where(mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0) + dreg_patch1_3_lon_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0) + dreg_patch1_3_lat_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0) # Case 1 - patch 2 - dreg_patch2_1_x = where(mask_case1, arrival_pts_2_x, 0.0) - dreg_patch2_1_y = where(mask_case1, arrival_pts_2_y, 0.0) - dreg_patch2_4_x = where(mask_case1, arrival_pts_2_x, 0.0) - dreg_patch2_4_y = where(mask_case1, arrival_pts_2_y, 0.0) - dreg_patch2_2_x = where(mask_case1, where(lvn_sys_pos, depart_pts_2_x, ps2_x), 0.0) - dreg_patch2_2_y = where(mask_case1, where(lvn_sys_pos, depart_pts_2_y, ps2_y), 0.0) - dreg_patch2_3_x = where(mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_x), 0.0) - dreg_patch2_3_y = where(mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_y), 0.0) + dreg_patch2_1_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0) + dreg_patch2_2_lat_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0) + dreg_patch2_3_lon_vmask = where(mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0) + dreg_patch2_3_lat_vmask = where(mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0) # ------------------------------------------------- Case 2a mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool # Case 2a - patch 0 - dreg_patch0_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch0_1_x) - dreg_patch0_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where( - mask_case2a, where(lvn_sys_pos, arrival_pts_2_x, ps1_x), dreg_patch0_2_x + dreg_patch0_1_lon_dsl = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case2a, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl ) - dreg_patch0_2_y = where( - mask_case2a, where(lvn_sys_pos, arrival_pts_2_y, ps1_y), dreg_patch0_2_y + dreg_patch0_2_lat_dsl = where( + mask_case2a, where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl ) - dreg_patch0_3_x = where(mask_case2a, depart_pts_2_x, dreg_patch0_3_x) - dreg_patch0_3_y = where(mask_case2a, depart_pts_2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where( - mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_x), dreg_patch0_4_x + dreg_patch0_3_lon_dsl = where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl ) - dreg_patch0_4_y = where( - mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_y), dreg_patch0_4_y + dreg_patch0_4_lat_dsl = where( + mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl ) # Case 2a - patch 1 - dreg_patch1_1_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_1_x) - dreg_patch1_1_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_1_y) - dreg_patch1_4_x = where(mask_case2a, arrival_pts_1_x, dreg_patch1_4_x) - dreg_patch1_4_y = where(mask_case2a, arrival_pts_1_y, dreg_patch1_4_y) - dreg_patch1_2_x = where( - mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_x), dreg_patch1_2_x + dreg_patch1_1_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_2_lon_vmask = where( + mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask ) - dreg_patch1_2_y = where( - mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_y), dreg_patch1_2_y + dreg_patch1_2_lat_vmask = where( + mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask ) - dreg_patch1_3_x = where( - mask_case2a, where(lvn_sys_pos, depart_pts_1_x, ps1_x), dreg_patch1_3_x + dreg_patch1_3_lon_vmask = where( + mask_case2a, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask ) - dreg_patch1_3_y = where( - mask_case2a, where(lvn_sys_pos, depart_pts_1_y, ps1_y), dreg_patch1_3_y + dreg_patch1_3_lat_vmask = where( + mask_case2a, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask ) # Case 2a - patch 2 - dreg_patch2_1_x = where(mask_case2a, 0.0, dreg_patch2_1_x) - dreg_patch2_1_y = where(mask_case2a, 0.0, dreg_patch2_1_y) - dreg_patch2_2_x = where(mask_case2a, 0.0, dreg_patch2_2_x) - dreg_patch2_2_y = where(mask_case2a, 0.0, dreg_patch2_2_y) - dreg_patch2_3_x = where(mask_case2a, 0.0, dreg_patch2_3_x) - dreg_patch2_3_y = where(mask_case2a, 0.0, dreg_patch2_3_y) - dreg_patch2_4_x = where(mask_case2a, 0.0, dreg_patch2_4_x) - dreg_patch2_4_y = where(mask_case2a, 0.0, dreg_patch2_4_y) + dreg_patch2_1_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) # -------------------------------------------------- Case 2b mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool # Case 2b - patch 0 - dreg_patch0_1_x = where(mask_case2b, arrival_pts_1_x, dreg_patch0_1_x) - dreg_patch0_1_y = where(mask_case2b, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where( + dreg_patch0_1_lon_dsl = where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( mask_case2b, - where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), - dreg_patch0_2_x, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_y = where( + dreg_patch0_2_lat_dsl = where( mask_case2b, - where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), - dreg_patch0_2_y, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_x = where(mask_case2b, ps2_x, dreg_patch0_3_x) - dreg_patch0_3_y = where(mask_case2b, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where( + dreg_patch0_3_lon_dsl = where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( mask_case2b, - where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), - dreg_patch0_4_x, + where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_y = where( + dreg_patch0_4_lat_dsl = where( mask_case2b, - where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), - dreg_patch0_4_y, + where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 2b - patch 1 - dreg_patch1_1_x = where(mask_case2b, 0.0, dreg_patch1_1_x) - dreg_patch1_1_y = where(mask_case2b, 0.0, dreg_patch1_1_y) - dreg_patch1_2_x = where(mask_case2b, 0.0, dreg_patch1_2_x) - dreg_patch1_2_y = where(mask_case2b, 0.0, dreg_patch1_2_y) - dreg_patch1_3_x = where(mask_case2b, 0.0, dreg_patch1_3_x) - dreg_patch1_3_y = where(mask_case2b, 0.0, dreg_patch1_3_y) - dreg_patch1_4_x = where(mask_case2b, 0.0, dreg_patch1_4_x) - dreg_patch1_4_y = where(mask_case2b, 0.0, dreg_patch1_4_y) + dreg_patch1_1_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) # Case 2b - patch 2 - dreg_patch2_1_x = where(mask_case2b, arrival_pts_2_x, dreg_patch2_1_x) - dreg_patch2_1_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_1_y) - dreg_patch2_4_x = where(mask_case2b, arrival_pts_2_x, dreg_patch2_4_x) - dreg_patch2_4_y = where(mask_case2b, arrival_pts_2_y, dreg_patch2_4_y) - dreg_patch2_2_x = where( - mask_case2b, where(lvn_sys_pos, depart_pts_2_x, ps2_x), dreg_patch2_2_x + dreg_patch2_1_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_2_lon_vmask = where( + mask_case2b, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask ) - dreg_patch2_2_y = where( - mask_case2b, where(lvn_sys_pos, depart_pts_2_y, ps2_y), dreg_patch2_2_y + dreg_patch2_2_lat_vmask = where( + mask_case2b, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask ) - dreg_patch2_3_x = where( - mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_x), dreg_patch2_3_x + dreg_patch2_3_lon_vmask = where( + mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask ) - dreg_patch2_3_y = where( - mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_y), dreg_patch2_3_y + dreg_patch2_3_lat_vmask = where( + mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask ) # flux area edge 1 and 2 - fl_e1_p1_lon = arrival_pts_1_x - fl_e1_p1_lat = arrival_pts_1_y - fl_e1_p2_lon = depart_pts_1_x - fl_e1_p2_lat = depart_pts_1_y - fl_e2_p1_lon = arrival_pts_2_x - fl_e2_p1_lat = arrival_pts_2_y - fl_e2_p2_lon = depart_pts_2_x - fl_e2_p2_lat = depart_pts_2_y + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl # ----------------------------------------------- Case 3a # Check whether flux area edge 2 intersects with triangle edge 1 @@ -427,56 +427,56 @@ def _divide_flux_area_list_stencil_01( tri_line1_p2_lat, ) # Case 3a - patch 0 - dreg_patch0_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch0_1_x) - dreg_patch0_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_2_x = where( + dreg_patch0_1_lon_dsl = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( mask_case3a, - where(lvn_sys_pos, arrival_pts_2_x, depart_pts_1_x), - dreg_patch0_2_x, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_y = where( + dreg_patch0_2_lat_dsl = where( mask_case3a, - where(lvn_sys_pos, arrival_pts_2_y, depart_pts_1_y), - dreg_patch0_2_y, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_x = where(mask_case3a, ps2_x, dreg_patch0_3_x) - dreg_patch0_3_y = where(mask_case3a, ps2_y, dreg_patch0_3_y) - dreg_patch0_4_x = where( + dreg_patch0_3_lon_dsl = where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( mask_case3a, - where(lvn_sys_pos, depart_pts_1_x, arrival_pts_2_x), - dreg_patch0_4_x, + where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_y = where( + dreg_patch0_4_lat_dsl = where( mask_case3a, - where(lvn_sys_pos, depart_pts_1_y, arrival_pts_2_y), - dreg_patch0_4_y, + where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 3a - patch 1 - dreg_patch1_1_x = where(mask_case3a, arrival_pts_1_x, dreg_patch1_1_x) - dreg_patch1_1_y = where(mask_case3a, arrival_pts_1_y, dreg_patch1_1_y) - dreg_patch1_2_x = where( - mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_x), dreg_patch1_2_x + dreg_patch1_1_lon_vmask = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where( + mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask ) - dreg_patch1_2_y = where( - mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_y), dreg_patch1_2_y + dreg_patch1_2_lat_vmask = where( + mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask ) - dreg_patch1_3_x = where(mask_case3a, depart_pts_1_x, dreg_patch1_3_x) - dreg_patch1_3_y = where(mask_case3a, depart_pts_1_y, dreg_patch1_3_y) - dreg_patch1_4_x = where( - mask_case3a, where(lvn_sys_pos, depart_pts_1_x, pi1_x), dreg_patch1_4_x + dreg_patch1_3_lon_vmask = where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where( + mask_case3a, where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask ) - dreg_patch1_4_y = where( - mask_case3a, where(lvn_sys_pos, depart_pts_1_y, pi1_y), dreg_patch1_4_y + dreg_patch1_4_lat_vmask = where( + mask_case3a, where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask ) # Case 3a - patch 2 - dreg_patch2_1_x = where(mask_case3a, 0.0, dreg_patch2_1_x) - dreg_patch2_1_y = where(mask_case3a, 0.0, dreg_patch2_1_y) - dreg_patch2_2_x = where(mask_case3a, 0.0, dreg_patch2_2_x) - dreg_patch2_2_y = where(mask_case3a, 0.0, dreg_patch2_2_y) - dreg_patch2_3_x = where(mask_case3a, 0.0, dreg_patch2_3_x) - dreg_patch2_3_y = where(mask_case3a, 0.0, dreg_patch2_3_y) - dreg_patch2_4_x = where(mask_case3a, 0.0, dreg_patch2_4_x) - dreg_patch2_4_y = where(mask_case3a, 0.0, dreg_patch2_4_y) + dreg_patch2_1_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) # ------------------------------------------------ Case 3b # Check whether flux area edge 1 intersects with triangle edge 2 @@ -502,47 +502,47 @@ def _divide_flux_area_list_stencil_01( tri_line2_p2_lat, ) # Case 3b - patch 0 - dreg_patch0_1_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_1_x) - dreg_patch0_1_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_1_y) - dreg_patch0_4_x = where(mask_case3b, arrival_pts_1_x, dreg_patch0_4_x) - dreg_patch0_4_y = where(mask_case3b, arrival_pts_1_y, dreg_patch0_4_y) - dreg_patch0_2_x = where( - mask_case3b, where(lvn_sys_pos, arrival_pts_2_x, pi2_x), dreg_patch0_2_x + dreg_patch0_1_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case3b, where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl ) - dreg_patch0_2_y = where( - mask_case3b, where(lvn_sys_pos, arrival_pts_2_y, pi2_y), dreg_patch0_2_y + dreg_patch0_2_lat_dsl = where( + mask_case3b, where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl ) - dreg_patch0_3_x = where( - mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_x), dreg_patch0_3_x + dreg_patch0_3_lon_dsl = where( + mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl ) - dreg_patch0_3_y = where( - mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_y), dreg_patch0_3_y + dreg_patch0_3_lat_dsl = where( + mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl ) # Case 3b - patch 1 - dreg_patch1_1_x = where(mask_case3b, 0.0, dreg_patch1_1_x) - dreg_patch1_1_y = where(mask_case3b, 0.0, dreg_patch1_1_y) - dreg_patch1_2_x = where(mask_case3b, 0.0, dreg_patch1_2_x) - dreg_patch1_2_y = where(mask_case3b, 0.0, dreg_patch1_2_y) - dreg_patch1_3_x = where(mask_case3b, 0.0, dreg_patch1_3_x) - dreg_patch1_3_y = where(mask_case3b, 0.0, dreg_patch1_3_y) - dreg_patch1_4_x = where(mask_case3b, 0.0, dreg_patch1_4_x) - dreg_patch1_4_y = where(mask_case3b, 0.0, dreg_patch1_4_y) + dreg_patch1_1_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 - dreg_patch2_1_x = where(mask_case3b, arrival_pts_2_x, dreg_patch2_1_x) - dreg_patch2_1_y = where(mask_case3b, arrival_pts_2_y, dreg_patch2_1_y) - dreg_patch2_2_x = where( - mask_case3b, where(lvn_sys_pos, depart_pts_2_x, pi2_x), dreg_patch2_2_x + dreg_patch2_1_lon_vmask = where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where( + mask_case3b, where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask ) - dreg_patch2_2_y = where( - mask_case3b, where(lvn_sys_pos, depart_pts_2_y, pi2_y), dreg_patch2_2_y + dreg_patch2_2_lat_vmask = where( + mask_case3b, where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask ) - dreg_patch2_3_x = where(mask_case3b, depart_pts_1_x, dreg_patch2_3_x) - dreg_patch2_3_y = where(mask_case3b, depart_pts_1_y, dreg_patch2_3_y) - dreg_patch2_4_x = where( - mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_x), dreg_patch2_4_x + dreg_patch2_3_lon_vmask = where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where( + mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask ) - dreg_patch2_4_y = where( - mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_y), dreg_patch2_4_y + dreg_patch2_4_lat_vmask = where( + mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask ) # --------------------------------------------- Case 4 @@ -555,49 +555,49 @@ def _divide_flux_area_list_stencil_01( mask_case4 = famask_bool & (not indices_previously_matched) # Case 4 - patch 0 - no change # Case 4 - patch 1 - dreg_patch1_1_x = where(mask_case4, 0.0, dreg_patch1_1_x) - dreg_patch1_1_y = where(mask_case4, 0.0, dreg_patch1_1_y) - dreg_patch1_2_x = where(mask_case4, 0.0, dreg_patch1_2_x) - dreg_patch1_2_y = where(mask_case4, 0.0, dreg_patch1_2_y) - dreg_patch1_3_x = where(mask_case4, 0.0, dreg_patch1_3_x) - dreg_patch1_3_y = where(mask_case4, 0.0, dreg_patch1_3_y) - dreg_patch1_4_x = where(mask_case4, 0.0, dreg_patch1_4_x) - dreg_patch1_4_y = where(mask_case4, 0.0, dreg_patch1_4_y) + dreg_patch1_1_lon_vmask = where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) # Case 4 - patch 2 - dreg_patch2_1_x = where(mask_case4, 0.0, dreg_patch2_1_x) - dreg_patch2_1_y = where(mask_case4, 0.0, dreg_patch2_1_y) - dreg_patch2_2_x = where(mask_case4, 0.0, dreg_patch2_2_x) - dreg_patch2_2_y = where(mask_case4, 0.0, dreg_patch2_2_y) - dreg_patch2_3_x = where(mask_case4, 0.0, dreg_patch2_3_x) - dreg_patch2_3_y = where(mask_case4, 0.0, dreg_patch2_3_y) - dreg_patch2_4_x = where(mask_case4, 0.0, dreg_patch2_4_x) - dreg_patch2_4_y = where(mask_case4, 0.0, dreg_patch2_4_y) + dreg_patch2_1_lon_vmask = where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case4, 0.0, dreg_patch2_4_lat_vmask) return ( - dreg_patch0_1_x, - dreg_patch0_1_y, - dreg_patch0_2_x, - dreg_patch0_2_y, - dreg_patch0_3_x, - dreg_patch0_3_y, - dreg_patch0_4_x, - dreg_patch0_4_y, - dreg_patch1_1_x, - dreg_patch1_1_y, - dreg_patch1_2_x, - dreg_patch1_2_y, - dreg_patch1_3_x, - dreg_patch1_3_y, - dreg_patch1_4_x, - dreg_patch1_4_y, - dreg_patch2_1_x, - dreg_patch2_1_y, - dreg_patch2_2_x, - dreg_patch2_2_y, - dreg_patch2_3_x, - dreg_patch2_3_y, - dreg_patch2_4_x, - dreg_patch2_4_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, ) @@ -608,30 +608,30 @@ def divide_flux_area_list_stencil_01( ptr_v3_lon: Field[[ECDim], float], ptr_v3_lat: Field[[ECDim], float], tangent_orientation_dsl: Field[[EdgeDim], float], - dreg_patch0_1_x: Field[[EdgeDim, KDim], float], - dreg_patch0_1_y: Field[[EdgeDim, KDim], float], - dreg_patch0_2_x: Field[[EdgeDim, KDim], float], - dreg_patch0_2_y: Field[[EdgeDim, KDim], float], - dreg_patch0_3_x: Field[[EdgeDim, KDim], float], - dreg_patch0_3_y: Field[[EdgeDim, KDim], float], - dreg_patch0_4_x: Field[[EdgeDim, KDim], float], - dreg_patch0_4_y: Field[[EdgeDim, KDim], float], - dreg_patch1_1_x: Field[[EdgeDim, KDim], float], - dreg_patch1_1_y: Field[[EdgeDim, KDim], float], - dreg_patch1_2_x: Field[[EdgeDim, KDim], float], - dreg_patch1_2_y: Field[[EdgeDim, KDim], float], - dreg_patch1_3_x: Field[[EdgeDim, KDim], float], - dreg_patch1_3_y: Field[[EdgeDim, KDim], float], - dreg_patch1_4_x: Field[[EdgeDim, KDim], float], - dreg_patch1_4_y: Field[[EdgeDim, KDim], float], - dreg_patch2_1_x: Field[[EdgeDim, KDim], float], - dreg_patch2_1_y: Field[[EdgeDim, KDim], float], - dreg_patch2_2_x: Field[[EdgeDim, KDim], float], - dreg_patch2_2_y: Field[[EdgeDim, KDim], float], - dreg_patch2_3_x: Field[[EdgeDim, KDim], float], - dreg_patch2_3_y: Field[[EdgeDim, KDim], float], - dreg_patch2_4_x: Field[[EdgeDim, KDim], float], - dreg_patch2_4_y: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], ): _divide_flux_area_list_stencil_01( famask_int, @@ -639,38 +639,38 @@ def divide_flux_area_list_stencil_01( ptr_v3_lon, ptr_v3_lat, tangent_orientation_dsl, - dreg_patch0_1_x, - dreg_patch0_1_y, - dreg_patch0_2_x, - dreg_patch0_2_y, - dreg_patch0_3_x, - dreg_patch0_3_y, - dreg_patch0_4_x, - dreg_patch0_4_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, out=( - dreg_patch0_1_x, - dreg_patch0_1_y, - dreg_patch0_2_x, - dreg_patch0_2_y, - dreg_patch0_3_x, - dreg_patch0_3_y, - dreg_patch0_4_x, - dreg_patch0_4_y, - dreg_patch1_1_x, - dreg_patch1_1_y, - dreg_patch1_2_x, - dreg_patch1_2_y, - dreg_patch1_3_x, - dreg_patch1_3_y, - dreg_patch1_4_x, - dreg_patch1_4_y, - dreg_patch2_1_x, - dreg_patch2_1_y, - dreg_patch2_2_x, - dreg_patch2_2_y, - dreg_patch2_3_x, - dreg_patch2_3_y, - dreg_patch2_4_x, - dreg_patch2_4_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, ), ) diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py index 8df569b509..dde713b651 100644 --- a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py +++ b/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py @@ -38,22 +38,22 @@ def _divide_flux_area_list_stencil_02( butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], - dreg_patch1_1_x: Field[[EdgeDim, KDim], float], - dreg_patch1_1_y: Field[[EdgeDim, KDim], float], - dreg_patch1_2_x: Field[[EdgeDim, KDim], float], - dreg_patch1_2_y: Field[[EdgeDim, KDim], float], - dreg_patch1_3_x: Field[[EdgeDim, KDim], float], - dreg_patch1_3_y: Field[[EdgeDim, KDim], float], - dreg_patch1_4_x: Field[[EdgeDim, KDim], float], - dreg_patch1_4_y: Field[[EdgeDim, KDim], float], - dreg_patch2_1_x: Field[[EdgeDim, KDim], float], - dreg_patch2_1_y: Field[[EdgeDim, KDim], float], - dreg_patch2_2_x: Field[[EdgeDim, KDim], float], - dreg_patch2_2_y: Field[[EdgeDim, KDim], float], - dreg_patch2_3_x: Field[[EdgeDim, KDim], float], - dreg_patch2_3_y: Field[[EdgeDim, KDim], float], - dreg_patch2_4_x: Field[[EdgeDim, KDim], float], - dreg_patch2_4_y: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], ) -> tuple[ Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], @@ -101,23 +101,23 @@ def _divide_flux_area_list_stencil_02( ) # patch1 in translated system - dreg_patch1_1_x = dreg_patch1_1_x - bf_cc_patch1_lon - dreg_patch1_1_y = dreg_patch1_1_y - bf_cc_patch1_lat - dreg_patch1_2_x = dreg_patch1_2_x - bf_cc_patch1_lon - dreg_patch1_2_y = dreg_patch1_2_y - bf_cc_patch1_lat - dreg_patch1_3_x = dreg_patch1_3_x - bf_cc_patch1_lon - dreg_patch1_3_y = dreg_patch1_3_y - bf_cc_patch1_lat - dreg_patch1_4_x = dreg_patch1_4_x - bf_cc_patch1_lon - dreg_patch1_4_y = dreg_patch1_4_y - bf_cc_patch1_lat + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat # patch2 in translated system - dreg_patch2_1_x = dreg_patch2_1_x - bf_cc_patch2_lon - dreg_patch2_1_y = dreg_patch2_1_y - bf_cc_patch2_lat - dreg_patch2_2_x = dreg_patch2_2_x - bf_cc_patch2_lon - dreg_patch2_2_y = dreg_patch2_2_y - bf_cc_patch2_lat - dreg_patch2_3_x = dreg_patch2_3_x - bf_cc_patch2_lon - dreg_patch2_3_y = dreg_patch2_3_y - bf_cc_patch2_lat - dreg_patch2_4_x = dreg_patch2_4_x - bf_cc_patch2_lon - dreg_patch2_4_y = dreg_patch2_4_y - bf_cc_patch2_lat + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat # Store global index of the underlying grid cell # Adapt dimensions to fit ofr multiple levels @@ -145,48 +145,48 @@ def _divide_flux_area_list_stencil_02( butterfly_blk_patch2_vnneg_3d = broadcast( butterfly_blk_patch2_vnneg, (EdgeDim, KDim) ) - patch1_cell_idx_dsl = where( + patch1_cell_idx_vmask = where( famask_bool, where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), int32(0), ) - patch2_cell_idx_dsl = where( + patch2_cell_idx_vmask = where( famask_bool, where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), int32(0), ) - patch1_cell_blk_dsl = where( + patch1_cell_blk_vmask = where( famask_bool, where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), int32(0), ) - patch2_cell_blk_dsl = where( + patch2_cell_blk_vmask = where( famask_bool, where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), int32(0), ) return ( - dreg_patch1_1_x, - dreg_patch1_1_y, - dreg_patch1_2_x, - dreg_patch1_2_y, - dreg_patch1_3_x, - dreg_patch1_3_y, - dreg_patch1_4_x, - dreg_patch1_4_y, - dreg_patch2_1_x, - dreg_patch2_1_y, - dreg_patch2_2_x, - dreg_patch2_2_y, - dreg_patch2_3_x, - dreg_patch2_3_y, - dreg_patch2_4_x, - dreg_patch2_4_y, - patch1_cell_idx_dsl, - patch1_cell_blk_dsl, - patch2_cell_idx_dsl, - patch2_cell_blk_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, ) @@ -206,26 +206,26 @@ def divide_flux_area_list_stencil_02( butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], - dreg_patch1_1_x: Field[[EdgeDim, KDim], float], - dreg_patch1_1_y: Field[[EdgeDim, KDim], float], - dreg_patch1_2_x: Field[[EdgeDim, KDim], float], - dreg_patch1_2_y: Field[[EdgeDim, KDim], float], - dreg_patch1_3_x: Field[[EdgeDim, KDim], float], - dreg_patch1_3_y: Field[[EdgeDim, KDim], float], - dreg_patch1_4_x: Field[[EdgeDim, KDim], float], - dreg_patch1_4_y: Field[[EdgeDim, KDim], float], - dreg_patch2_1_x: Field[[EdgeDim, KDim], float], - dreg_patch2_1_y: Field[[EdgeDim, KDim], float], - dreg_patch2_2_x: Field[[EdgeDim, KDim], float], - dreg_patch2_2_y: Field[[EdgeDim, KDim], float], - dreg_patch2_3_x: Field[[EdgeDim, KDim], float], - dreg_patch2_3_y: Field[[EdgeDim, KDim], float], - dreg_patch2_4_x: Field[[EdgeDim, KDim], float], - dreg_patch2_4_y: Field[[EdgeDim, KDim], float], - patch1_cell_idx_dsl: Field[[EdgeDim, KDim], int32], - patch1_cell_blk_dsl: Field[[EdgeDim, KDim], int32], - patch2_cell_idx_dsl: Field[[EdgeDim, KDim], int32], - patch2_cell_blk_dsl: Field[[EdgeDim, KDim], int32], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], + patch1_cell_idx_vmask: Field[[EdgeDim, KDim], int32], + patch1_cell_blk_vmask: Field[[EdgeDim, KDim], int32], + patch2_cell_idx_vmask: Field[[EdgeDim, KDim], int32], + patch2_cell_blk_vmask: Field[[EdgeDim, KDim], int32], ): _divide_flux_area_list_stencil_02( famask_int, @@ -242,42 +242,42 @@ def divide_flux_area_list_stencil_02( butterfly_idx_patch2_vnneg, butterfly_blk_patch2_vnpos, butterfly_blk_patch2_vnneg, - dreg_patch1_1_x, - dreg_patch1_1_y, - dreg_patch1_2_x, - dreg_patch1_2_y, - dreg_patch1_3_x, - dreg_patch1_3_y, - dreg_patch1_4_x, - dreg_patch1_4_y, - dreg_patch2_1_x, - dreg_patch2_1_y, - dreg_patch2_2_x, - dreg_patch2_2_y, - dreg_patch2_3_x, - dreg_patch2_3_y, - dreg_patch2_4_x, - dreg_patch2_4_y, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, out=( - dreg_patch1_1_x, - dreg_patch1_1_y, - dreg_patch1_2_x, - dreg_patch1_2_y, - dreg_patch1_3_x, - dreg_patch1_3_y, - dreg_patch1_4_x, - dreg_patch1_4_y, - dreg_patch2_1_x, - dreg_patch2_1_y, - dreg_patch2_2_x, - dreg_patch2_2_y, - dreg_patch2_3_x, - dreg_patch2_3_y, - dreg_patch2_4_x, - dreg_patch2_4_y, - patch1_cell_idx_dsl, - patch1_cell_blk_dsl, - patch2_cell_idx_dsl, - patch2_cell_blk_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, ), ) From 145b159cf2eae327ff4763dfa97c18f96635c03e Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 27 Jul 2023 10:39:08 +0200 Subject: [PATCH 081/105] Adapt+add stencils+tests for advection scheme 52 --- .../advection/btraj_dreg_stencil_02.py | 4 +- .../advection/btraj_dreg_stencil_03.py | 3 - .../advection/hflx_limiter_mo_stencil_01a.py | 14 +- .../advection/hflx_limiter_mo_stencil_01b.py | 29 +- .../prep_gauss_quadrature_c_list_stencil.py | 21 +- .../upwind_hflux_miura3_stencil_01.py | 115 +++ .../upwind_hflux_miura_cycl_stencil_01.py | 2 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 3 - .../upwind_hflux_miura_stencil_02.py | 55 -- advection/tests/test_btraj_dreg_stencil_01.py | 67 ++ advection/tests/test_btraj_dreg_stencil_02.py | 74 ++ advection/tests/test_btraj_dreg_stencil_03.py | 260 +++++++ .../test_divide_flux_area_list_stencil_01.py | 712 ++++++++++++++++++ .../test_divide_flux_area_list_stencil_02.py | 333 ++++++++ .../test_hflux_ffsl_hybrid_stencil_01a.py | 156 ++++ .../test_hflux_ffsl_hybrid_stencil_02.py | 55 ++ .../tests/test_hflx_limiter_mo_stencil_01a.py | 71 ++ .../tests/test_hflx_limiter_mo_stencil_01b.py | 121 +++ ...st_prep_gauss_quadrature_c_list_stencil.py | 504 +++++++++++++ .../test_prep_gauss_quadrature_c_stencil.py | 464 ++++++++++++ .../tests/test_recon_lsq_cell_c_stencil.py | 683 +++++++++++++++++ .../test_recon_lsq_cell_c_svd_stencil.py | 337 +++++++++ ...y => test_recon_lsq_cell_l_svd_stencil.py} | 12 +- .../test_upwind_hflux_miura3_stencil_01.py | 165 ++++ ...test_upwind_hflux_miura_cycl_stencil_01.py | 82 ++ ...test_upwind_hflux_miura_cycl_stencil_02.py | 102 +++ ...est_upwind_hflux_miura_cycl_stencil_03a.py | 47 ++ ...est_upwind_hflux_miura_cycl_stencil_03b.py | 53 ++ .../test_upwind_hflux_miura_stencil_01.py | 60 +- advection/tests/test_utils/simple_mesh.py | 29 + 30 files changed, 4499 insertions(+), 134 deletions(-) create mode 100644 advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py delete mode 100644 advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py create mode 100644 advection/tests/test_btraj_dreg_stencil_01.py create mode 100644 advection/tests/test_btraj_dreg_stencil_02.py create mode 100644 advection/tests/test_btraj_dreg_stencil_03.py create mode 100644 advection/tests/test_divide_flux_area_list_stencil_01.py create mode 100644 advection/tests/test_divide_flux_area_list_stencil_02.py create mode 100644 advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py create mode 100644 advection/tests/test_hflux_ffsl_hybrid_stencil_02.py create mode 100644 advection/tests/test_hflx_limiter_mo_stencil_01a.py create mode 100644 advection/tests/test_hflx_limiter_mo_stencil_01b.py create mode 100644 advection/tests/test_prep_gauss_quadrature_c_list_stencil.py create mode 100644 advection/tests/test_prep_gauss_quadrature_c_stencil.py create mode 100644 advection/tests/test_recon_lsq_cell_c_stencil.py create mode 100644 advection/tests/test_recon_lsq_cell_c_svd_stencil.py rename advection/tests/{test_upwind_hflux_miura_stencil_02.py => test_recon_lsq_cell_l_svd_stencil.py} (89%) create mode 100644 advection/tests/test_upwind_hflux_miura3_stencil_01.py create mode 100644 advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py create mode 100644 advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py create mode 100644 advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py create mode 100644 advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py index 1c4d116503..cd734f47f3 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_02.py @@ -29,7 +29,7 @@ def _btraj_dreg_stencil_02( traj_length = sqrt(p_vn**2 + p_vt**2) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) opt_famask_dsl = where( - traj_length > 0.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) + traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) ) return opt_famask_dsl @@ -40,7 +40,7 @@ def btraj_dreg_stencil_02( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], edge_cell_length: Field[[ECDim], float], - opt_famask_dsl: Field[[EdgeDim, KDim], int32], p_dt: float, + opt_famask_dsl: Field[[EdgeDim, KDim], int32], ): _btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=opt_famask_dsl) diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py index 30283e3df3..1df7ea7daa 100644 --- a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py +++ b/advection/src/icon4py/advection/btraj_dreg_stencil_03.py @@ -50,9 +50,6 @@ def _btraj_dreg_stencil_03( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: - # this is comented out: lvn_sys_pos = where(lvn_sys_pos_int == int32(1), broadcast(True, (EdgeDim, KDim)), broadcast(False, (EdgeDim, KDim))) - # this is comented out: lvn_sys_pos = where(lvn_sys_pos_int == int32(1), True, False) - lvn_pos = where(p_vn >= 0.0, True, False) p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py index cf86fbce4f..97dd9ead47 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py @@ -21,34 +21,30 @@ def _hflx_limiter_mo_stencil_01a( p_mflx_tracer_h: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], - z_mflx_low: Field[[EdgeDim, KDim], float], - z_anti: Field[[EdgeDim, KDim], float], p_cc: Field[[CellDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: z_mflx_low = 0.5 * ( p_mass_flx_e * (p_cc(E2C[0]) + p_cc(E2C[1])) - - abs(p_mass_flx_e) * (p_cc(E2C[0]) - p_cc(E2C[1])) + - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) ) - # this is comented out: z_mflx_low = 0.5 * ( p_mass_flx_e * ( p_cc(E2C[0]) + p_cc(E2C[1]) ) ) + z_anti = p_mflx_tracer_h - z_mflx_low - return (z_anti, z_mflx_low) + return (z_mflx_low, z_anti) @program def hflx_limiter_mo_stencil_01a( p_mflx_tracer_h: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], z_mflx_low: Field[[EdgeDim, KDim], float], z_anti: Field[[EdgeDim, KDim], float], - p_cc: Field[[CellDim, KDim], float], ): _hflx_limiter_mo_stencil_01a( p_mflx_tracer_h, p_mass_flx_e, - z_mflx_low, - z_anti, p_cc, - out=(z_anti, z_mflx_low), + out=(z_mflx_low, z_anti), ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py index ada965d084..395ce606b9 100644 --- a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py @@ -20,18 +20,17 @@ neighbor_sum, ) -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2E, C2CE, CEDim, C2EDim, CellDim, EdgeDim, KDim @field_operator def _hflx_limiter_mo_stencil_01b( - geofac_div: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CEDim], float], p_rhodz_now: Field[[CellDim, KDim], float], p_rhodz_new: Field[[CellDim, KDim], float], z_mflx_low: Field[[EdgeDim, KDim], float], z_anti: Field[[EdgeDim, KDim], float], p_cc: Field[[CellDim, KDim], float], - z_fluxdiv_c: Field[[CellDim, KDim], float], p_dtime: float, ) -> tuple[ Field[[CellDim, KDim], float], @@ -40,12 +39,24 @@ def _hflx_limiter_mo_stencil_01b( Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], ]: + zero = broadcast(0.0, (CellDim, KDim)) - z_mflx_anti = broadcast(p_rhodz_new, (CellDim, C2EDim, KDim)) + z_mflx_anti_1 = ( p_dtime * geofac_div(C2CE[0]) / p_rhodz_new + * z_anti(C2E[0]) ) + z_mflx_anti_2 = ( p_dtime * geofac_div(C2CE[1]) / p_rhodz_new + * z_anti(C2E[1]) ) + z_mflx_anti_3 = ( p_dtime * geofac_div(C2CE[2]) / p_rhodz_new + * z_anti(C2E[2]) ) - z_mflx_anti_in = -1.0 * neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) - z_mflx_anti_out = neighbor_sum(p_rhodz_now(C2E) * z_mflx_anti, axis=C2EDim) - z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div, axis=C2EDim) + z_mflx_anti_in = -1.0 * ( minimum(zero, z_mflx_anti_1) + + minimum(zero, z_mflx_anti_2) + + minimum(zero, z_mflx_anti_3) ) + + z_mflx_anti_out = ( maximum(zero, z_mflx_anti_1) + + maximum(zero, z_mflx_anti_2) + + maximum(zero, z_mflx_anti_3) ) + + z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div(C2CE), axis=C2EDim) z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new z_tracer_max = maximum(p_cc, z_tracer_new_low) @@ -62,13 +73,12 @@ def _hflx_limiter_mo_stencil_01b( @program def hflx_limiter_mo_stencil_01b( - geofac_div: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CEDim], float], p_rhodz_now: Field[[CellDim, KDim], float], p_rhodz_new: Field[[CellDim, KDim], float], z_mflx_low: Field[[EdgeDim, KDim], float], z_anti: Field[[EdgeDim, KDim], float], p_cc: Field[[CellDim, KDim], float], - z_fluxdiv_c: Field[[CellDim, KDim], float], p_dtime: float, z_mflx_anti_in: Field[[CellDim, KDim], float], z_mflx_anti_out: Field[[CellDim, KDim], float], @@ -83,7 +93,6 @@ def hflx_limiter_mo_stencil_01b( z_mflx_low, z_anti, p_cc, - z_fluxdiv_c, p_dtime, out=( z_mflx_anti_in, diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py index 37119dcc76..213d84ea1d 100644 --- a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py @@ -58,7 +58,7 @@ def _prep_gauss_quadrature_c_list_stencil( wgt_eta_2: float, dbl_eps: float, eps: float, - p_dreg_area: Field[[EdgeDim, KDim], float], + p_dreg_area_in: Field[[EdgeDim, KDim], float], ) -> tuple[ Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], @@ -291,10 +291,10 @@ def _prep_gauss_quadrature_c_list_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_x**3 ) p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**3 - + wgt_t_detjac_2 * z_gauss_pts_2_y**3 - + wgt_t_detjac_3 * z_gauss_pts_3_y**3 - + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_9 = ( wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y @@ -309,10 +309,9 @@ def _prep_gauss_quadrature_c_list_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 ) - p_dreg_area = p_dreg_area + p_quad_vector_sum_1 + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 return ( - p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, @@ -323,6 +322,7 @@ def _prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, + p_dreg_area, ) @@ -367,7 +367,7 @@ def prep_gauss_quadrature_c_list_stencil( wgt_eta_2: float, dbl_eps: float, eps: float, - p_dreg_area: Field[[EdgeDim, KDim], float], + p_dreg_area_in: Field[[EdgeDim, KDim], float], p_quad_vector_sum_1: Field[[EdgeDim, KDim], float], p_quad_vector_sum_2: Field[[EdgeDim, KDim], float], p_quad_vector_sum_3: Field[[EdgeDim, KDim], float], @@ -378,6 +378,7 @@ def prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_8: Field[[EdgeDim, KDim], float], p_quad_vector_sum_9: Field[[EdgeDim, KDim], float], p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + p_dreg_area: Field[[EdgeDim, KDim], float], ): _prep_gauss_quadrature_c_list_stencil( famask_int, @@ -419,9 +420,8 @@ def prep_gauss_quadrature_c_list_stencil( wgt_eta_2, dbl_eps, eps, - p_dreg_area, + p_dreg_area_in, out=( - p_dreg_area, p_quad_vector_sum_1, p_quad_vector_sum_2, p_quad_vector_sum_3, @@ -432,5 +432,6 @@ def prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, + p_dreg_area, ), ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py b/advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py new file mode 100644 index 0000000000..a655fee634 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py @@ -0,0 +1,115 @@ +# 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.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C + + +@field_operator +def _upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e_miura3 = ( + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) * z_quad_vector_sum_1 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) * z_quad_vector_sum_2 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) * z_quad_vector_sum_3 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_4(E2C[1]), z_lsq_coeff_4(E2C[0])) * z_quad_vector_sum_4 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_5(E2C[1]), z_lsq_coeff_5(E2C[0])) * z_quad_vector_sum_5 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_6(E2C[1]), z_lsq_coeff_6(E2C[0])) * z_quad_vector_sum_6 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_7(E2C[1]), z_lsq_coeff_7(E2C[0])) * z_quad_vector_sum_7 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_8(E2C[1]), z_lsq_coeff_8(E2C[0])) * z_quad_vector_sum_8 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_9(E2C[1]), z_lsq_coeff_9(E2C[0])) * z_quad_vector_sum_9 + + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0])) * z_quad_vector_sum_10 + ) / z_dreg_area * p_mass_flx_e + + return p_out_e_miura3 + +@program +def upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_out_e_miura3: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum_1, + z_quad_vector_sum_2, + z_quad_vector_sum_3, + z_quad_vector_sum_4, + z_quad_vector_sum_5, + z_quad_vector_sum_6, + z_quad_vector_sum_7, + z_quad_vector_sum_8, + z_quad_vector_sum_9, + z_quad_vector_sum_10, + z_dreg_area, + p_mass_flx_e, + cell_rel_idx_dsl, + out=(p_out_e_miura3) + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py index 5466fb2a03..899000f283 100644 --- a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -29,7 +29,7 @@ def _upwind_hflux_miura_cycl_stencil_01( cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - z_tracer_mflx_dsl = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl(E2C[1]), z_lsq_coeff_1_dsl(E2C[0])) + z_tracer_mflx_dsl = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl(E2C[1]), z_lsq_coeff_1_dsl(E2C[0])) + distv_bary_1 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_dsl(E2C[1]), z_lsq_coeff_2_dsl(E2C[0])) + distv_bary_2 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_dsl(E2C[1]), z_lsq_coeff_3_dsl(E2C[0])) ) * p_mass_flx_e diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py index c654b484bd..4bb18f5302 100644 --- a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -20,7 +20,6 @@ @field_operator def _upwind_hflux_miura_cycl_stencil_03b( - p_ncycl: int32, z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], @@ -33,14 +32,12 @@ def _upwind_hflux_miura_cycl_stencil_03b( @program def upwind_hflux_miura_cycl_stencil_03b( - p_ncycl: int32, z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], p_out_e: Field[[EdgeDim, KDim], float], ): _upwind_hflux_miura_cycl_stencil_03b( - p_ncycl, z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, z_tracer_mflx_3_dsl, diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py deleted file mode 100644 index fb0c84d3f2..0000000000 --- a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py +++ /dev/null @@ -1,55 +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.next.common import Field -from gt4py.next.ffront.decorator import field_operator, program - -from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim - - -@field_operator -def _upwind_hflux_miura_stencil_02( - p_cc: Field[[CellDim, KDim], float], - lsq_pseudoinv_1: Field[[CECDim], float], - lsq_pseudoinv_2: Field[[CECDim], float], -) -> tuple[ - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], -]: - p_coeff_1 = p_cc - p_coeff_2 = ( - lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) - + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) - + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) - ) - p_coeff_3 = ( - lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) - + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) - + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) - ) - return p_coeff_1, p_coeff_2, p_coeff_3 - - -@program -def upwind_hflux_miura_stencil_02( - p_cc: Field[[CellDim, KDim], float], - lsq_pseudoinv_1: Field[[CECDim], float], - lsq_pseudoinv_2: Field[[CECDim], float], - p_coeff_1: Field[[CellDim, KDim], float], - p_coeff_2: Field[[CellDim, KDim], float], - p_coeff_3: Field[[CellDim, KDim], float], -): - _upwind_hflux_miura_stencil_02( - p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2, p_coeff_3) - ) diff --git a/advection/tests/test_btraj_dreg_stencil_01.py b/advection/tests/test_btraj_dreg_stencil_01.py new file mode 100644 index 0000000000..56ce5b035f --- /dev/null +++ b/advection/tests/test_btraj_dreg_stencil_01.py @@ -0,0 +1,67 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 +from icon4py.common.dimension import CellDim, EdgeDim, KDim + +from .test_utils.helpers import _shape, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_01_numpy( + lcounterclock: bool, + p_vn: np.array, + tangent_orientation: np.array, +): + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) + + lvn_sys_pos_true = np.where( tangent_orientation * p_vn >= 0.0, True, False) + + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) + + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) + + return lvn_sys_pos + + +def test_btraj_dreg_stencil_01(): + mesh = SimpleMesh() + lcounterclock = True + p_vn = random_field(mesh, EdgeDim, KDim) + + tangent_orientation = random_field(mesh, EdgeDim) + + lvn_sys_pos = zero_field(mesh, EdgeDim, KDim, dtype=bool) + + ref = btraj_dreg_stencil_01_numpy( + lcounterclock, + np.asarray(p_vn), + np.asarray(tangent_orientation), + ) + + btraj_dreg_stencil_01( + lcounterclock, + p_vn, + tangent_orientation, + lvn_sys_pos, + offset_provider={}, + + ) + + assert np.allclose(ref, lvn_sys_pos) diff --git a/advection/tests/test_btraj_dreg_stencil_02.py b/advection/tests/test_btraj_dreg_stencil_02.py new file mode 100644 index 0000000000..a7f2385898 --- /dev/null +++ b/advection/tests/test_btraj_dreg_stencil_02.py @@ -0,0 +1,74 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider +from icon4py.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 +from icon4py.common.dimension import E2CDim, ECDim, EdgeDim, KDim + +from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from .test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_02_numpy( + p_vn: np.array, + p_vt: np.array, + edge_cell_length: np.array, + p_dt: float, +) -> np.array: + lvn_pos = np.where(p_vn >= 0.0, True, False) + + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt + + edge_cell_length = np.expand_dims(edge_cell_length, axis=-1) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + + opt_famask_dsl = np.where( + traj_length > (0.25 * np.broadcast_to(e2c_length, p_vn.shape)), + np.int32(1), + np.int32(0) + ) + + return opt_famask_dsl + + +def test_btraj_dreg_stencil_02(): + mesh = SimpleMesh() + p_vn = random_field(mesh, EdgeDim, KDim) + p_vt = random_field(mesh, EdgeDim, KDim) + edge_cell_length = np.asarray(mesh.e2c, dtype=float) + edge_cell_length_new = as_1D_sparse_field(edge_cell_length, ECDim) + p_dt= 1.0 + opt_famask_dsl = zero_field(mesh, EdgeDim, KDim, dtype=int32) + + ref = btraj_dreg_stencil_02_numpy( + np.asarray(p_vn), + np.asarray(p_vt), + np.asarray(edge_cell_length), + p_dt + ) + + btraj_dreg_stencil_02( + p_vn, + p_vt, + edge_cell_length_new, + p_dt, + opt_famask_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + + assert np.allclose(ref, np.asarray(opt_famask_dsl)) diff --git a/advection/tests/test_btraj_dreg_stencil_03.py b/advection/tests/test_btraj_dreg_stencil_03.py new file mode 100644 index 0000000000..4e6e545604 --- /dev/null +++ b/advection/tests/test_btraj_dreg_stencil_03.py @@ -0,0 +1,260 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 +from icon4py.common.dimension import KDim, E2CDim, ECDim, EdgeDim, CellDim + +from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field, constant_field +from .test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_03_numpy( + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, + p_dt: float, +) -> tuple[np.array]: + lvn_pos = np.where(p_vn >= 0.0, True, False) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) + + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) + + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = ( + np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_2_y = ( + np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + ) + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = ( + np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_4_y = ( + np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + ) + + pn_cell_1 = np.where( + lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1] + ) + pn_cell_2 = np.where( + lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1] + ) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + + p_coords_dreg_v_1_lon_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + ) + p_coords_dreg_v_2_lon_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + ) + p_coords_dreg_v_3_lon_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + ) + p_coords_dreg_v_4_lon_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + ) + p_coords_dreg_v_1_lat_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + ) + p_coords_dreg_v_2_lat_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + ) + p_coords_dreg_v_3_lat_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + ) + p_coords_dreg_v_4_lat_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + ) + + return ( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + ) + + +def test_btraj_dreg_stencil_03(): + mesh = SimpleMesh() + + p_vn = random_field(mesh, EdgeDim, KDim) + p_vt = random_field(mesh, EdgeDim, KDim) + cell_idx = np.asarray(mesh.e2c, dtype=int32) + cell_idx_new = as_1D_sparse_field(cell_idx, ECDim) + cell_blk = constant_field(mesh, 1, EdgeDim, E2CDim, dtype=int32) + cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) + + edge_verts_1_x = random_field(mesh, EdgeDim) + edge_verts_2_x = random_field(mesh, EdgeDim) + edge_verts_1_y = random_field(mesh, EdgeDim) + edge_verts_2_y = random_field(mesh, EdgeDim) + pos_on_tplane_e_1_x = random_field(mesh, EdgeDim) + pos_on_tplane_e_2_x = random_field(mesh, EdgeDim) + pos_on_tplane_e_1_y = random_field(mesh, EdgeDim) + pos_on_tplane_e_2_y = random_field(mesh, EdgeDim) + primal_normal_cell_x = random_field(mesh, EdgeDim, E2CDim) + primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) + dual_normal_cell_x = random_field(mesh, EdgeDim, E2CDim) + dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) + primal_normal_cell_y = random_field(mesh, EdgeDim, E2CDim) + primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) + dual_normal_cell_y = random_field(mesh, EdgeDim, E2CDim) + dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) + lvn_sys_pos = constant_field(mesh, True, EdgeDim, KDim, dtype=bool) + p_dt = 2.0 + p_cell_idx = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_blk = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_lat_dsl = random_field(mesh, EdgeDim, KDim) + + ( + p_cell_idx_ref, + p_cell_rel_idx_dsl_ref, + p_cell_blk_ref, + p_coords_dreg_v_1_lon_dsl_ref, + p_coords_dreg_v_2_lon_dsl_ref, + p_coords_dreg_v_3_lon_dsl_ref, + p_coords_dreg_v_4_lon_dsl_ref, + p_coords_dreg_v_1_lat_dsl_ref, + p_coords_dreg_v_2_lat_dsl_ref, + p_coords_dreg_v_3_lat_dsl_ref, + p_coords_dreg_v_4_lat_dsl_ref, + + ) = btraj_dreg_stencil_03_numpy( + np.asarray(p_vn), + np.asarray(p_vt), + np.asarray(cell_idx), + np.asarray(cell_blk), + np.asarray(edge_verts_1_x), + np.asarray(edge_verts_2_x), + np.asarray(edge_verts_1_y), + np.asarray(edge_verts_2_y), + np.asarray(pos_on_tplane_e_1_x), + np.asarray(pos_on_tplane_e_2_x), + np.asarray(pos_on_tplane_e_1_y), + np.asarray(pos_on_tplane_e_2_y), + np.asarray(primal_normal_cell_x), + np.asarray(primal_normal_cell_y), + np.asarray(dual_normal_cell_x), + np.asarray(dual_normal_cell_y), + np.asarray(lvn_sys_pos), + p_dt, + ) + + + btraj_dreg_stencil_03( + p_vn, + p_vt, + cell_idx_new, + cell_blk_new, + edge_verts_1_x, + edge_verts_2_x, + edge_verts_1_y, + edge_verts_2_y, + pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y, + primal_normal_cell_x_new, + primal_normal_cell_y_new, + dual_normal_cell_x_new, + dual_normal_cell_y_new, + lvn_sys_pos, + p_dt, + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(p_cell_idx, p_cell_idx_ref) + assert np.allclose(p_cell_rel_idx_dsl, p_cell_rel_idx_dsl_ref) + assert np.allclose(p_cell_blk, p_cell_blk_ref) + assert np.allclose(p_coords_dreg_v_1_lon_dsl, p_coords_dreg_v_1_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_2_lon_dsl, p_coords_dreg_v_2_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_3_lon_dsl, p_coords_dreg_v_3_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_4_lon_dsl, p_coords_dreg_v_4_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_1_lat_dsl, p_coords_dreg_v_1_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_2_lat_dsl, p_coords_dreg_v_2_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_3_lat_dsl, p_coords_dreg_v_3_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_4_lat_dsl, p_coords_dreg_v_4_lat_dsl_ref) diff --git a/advection/tests/test_divide_flux_area_list_stencil_01.py b/advection/tests/test_divide_flux_area_list_stencil_01.py new file mode 100644 index 0000000000..af22e53458 --- /dev/null +++ b/advection/tests/test_divide_flux_area_list_stencil_01.py @@ -0,0 +1,712 @@ +# 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 + +import numpy as np + +from icon4py.advection.divide_flux_area_list_stencil_01 import ( + divide_flux_area_list_stencil_01, +) +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim + +from .test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from .test_utils.simple_mesh import SimpleMesh + +# FUNCTIONS +# Checking turn when travelling along three points, used to check whether lines inters. +def ccw( + p0_lon: np.array, + p0_lat: np.array, + p1_lon: np.array, + p1_lat: np.array, + p2_lon: np.array, + p2_lat: np.array, +) -> np.array: + + dx1 = p1_lon - p0_lon + dy1 = p1_lat - p0_lat + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + return ccw_out + +# Checks whether two lines intersect +def lintersect( + line1_p1_lon: np.array, + line1_p1_lat: np.array, + line1_p2_lon: np.array, + line1_p2_lat: np.array, + line2_p1_lon: np.array, + line2_p1_lat: np.array, + line2_p2_lon: np.array, + line2_p2_lat: np.array, +) -> np.array: + + intersect1 = ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + +# Compute intersection point of two lines in 2D +def line_intersect( + line1_p1_lon: np.array, + line1_p1_lat: np.array, + line1_p2_lon: np.array, + line1_p2_lat: np.array, + line2_p1_lon: np.array, + line2_p1_lat: np.array, + line2_p2_lon: np.array, + line2_p2_lat: np.array, +) -> tuple[np.array]: + + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + intersect_1 = ( + line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon + ) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + + +def divide_flux_area_list_stencil_01_numpy( + e2c: np.array, + famask_int: np.array, + p_vn: np.array, + ptr_v3_lon: np.array, + ptr_v3_lat: np.array, + tangent_orientation_dsl: np.array, + dreg_patch0_1_lon_dsl: np.array, + dreg_patch0_1_lat_dsl: np.array, + dreg_patch0_2_lon_dsl: np.array, + dreg_patch0_2_lat_dsl: np.array, + dreg_patch0_3_lon_dsl: np.array, + dreg_patch0_3_lat_dsl: np.array, + dreg_patch0_4_lon_dsl: np.array, + dreg_patch0_4_lat_dsl: np.array, +): + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) + + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl + + lvn_pos = np.where(p_vn >= 0.0, True, False) + + # get flux area departure-line segment + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl + + # get triangle edge 1 (A1V3) + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl + tri_line1_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line1_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + # get triangle edge 2 (A2V3) + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line2_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, True, False + ) + famask_bool = np.where(famask_int == int32(1), True, False) + # ------------------------------------------------- Case 1 + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + ps1_x, ps1_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + # Case 1 - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case1, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case1, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + ) + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + ) + # Case 1 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0) + dreg_patch1_2_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0) + dreg_patch1_3_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0) + dreg_patch1_3_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0) + # Case 1 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0) + dreg_patch2_2_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0) + dreg_patch2_3_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0) + dreg_patch2_3_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0) + + # ------------------------------------------------- Case 2a + mask_case2a = np.logical_and.reduce([lintersect_line1, np.logical_not(lintersect_line2), famask_bool]) + # Case 2a - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2a, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2a, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2a, np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2a, np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + ) + # Case 2a - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_2_lon_vmask = np.where( + mask_case2a, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case2a, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case2a, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case2a, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask + ) + # Case 2a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + + # -------------------------------------------------- Case 2b + mask_case2b = np.logical_and.reduce([lintersect_line2, np.logical_not(lintersect_line1), famask_bool]) + # Case 2b - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 2b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) + # Case 2b - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_2_lon_vmask = np.where( + mask_case2b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case2b, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case2b, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case2b, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask + ) + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl + + # ----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = lintersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) + pi1_x, pi1_y = line_intersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Case 3a - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 3a - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where( + mask_case3a, np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case3a, np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask + ) + dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where( + mask_case3a, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case3a, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask + ) + # Case 3a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + + # ------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = lintersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = line_intersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + # Case 3b - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3b, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3b, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case3b, np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case3b, np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl + ) + # Case 3b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + # Case 3b - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where( + mask_case3b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case3b, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask + ) + dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where( + mask_case3b, np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case3b, np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask + ) + + # --------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = ( + np.logical_or.reduce([mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1]) + ) + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = np.logical_and.reduce([famask_bool, np.logical_not(indices_previously_matched)]) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + # Case 4 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + +def test_divide_flux_area_list_stencil_01(): + mesh = SimpleMesh() + + famask_int = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_vn = random_field(mesh, EdgeDim, KDim) + ptr_v3_lon = random_field(mesh, EdgeDim, E2CDim) + ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) + ptr_v3_lat = random_field(mesh, EdgeDim, E2CDim) + ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) + tangent_orientation_dsl = random_field(mesh, EdgeDim) + dreg_patch0_1_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = zero_field(mesh, EdgeDim, KDim) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20, + ref_21, + ref_22, + ref_23, + ref_24) = divide_flux_area_list_stencil_01_numpy( + mesh.e2c, + np.asarray(famask_int), + np.asarray(p_vn), + np.asarray(ptr_v3_lon), + np.asarray(ptr_v3_lat), + np.asarray(tangent_orientation_dsl), + np.asarray(dreg_patch0_1_lon_dsl), + np.asarray(dreg_patch0_1_lat_dsl), + np.asarray(dreg_patch0_2_lon_dsl), + np.asarray(dreg_patch0_2_lat_dsl), + np.asarray(dreg_patch0_3_lon_dsl), + np.asarray(dreg_patch0_3_lat_dsl), + np.asarray(dreg_patch0_4_lon_dsl), + np.asarray(dreg_patch0_4_lat_dsl), + ) + + divide_flux_area_list_stencil_01( + famask_int, + p_vn, + ptr_v3_lon_field, + ptr_v3_lat_field, + tangent_orientation_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(dreg_patch0_1_lon_dsl, ref_1) + assert np.allclose(dreg_patch0_1_lat_dsl, ref_2) + assert np.allclose(dreg_patch0_2_lon_dsl, ref_3) + assert np.allclose(dreg_patch0_2_lat_dsl, ref_4) + assert np.allclose(dreg_patch0_3_lon_dsl, ref_5) + assert np.allclose(dreg_patch0_3_lat_dsl, ref_6) + assert np.allclose(dreg_patch0_4_lon_dsl, ref_7) + assert np.allclose(dreg_patch0_4_lat_dsl, ref_8) + assert np.allclose(dreg_patch1_1_lon_vmask, ref_9) + assert np.allclose(dreg_patch1_1_lat_vmask, ref_10) + assert np.allclose(dreg_patch1_2_lon_vmask, ref_11) + assert np.allclose(dreg_patch1_2_lat_vmask, ref_12) + assert np.allclose(dreg_patch1_3_lon_vmask, ref_13) + assert np.allclose(dreg_patch1_3_lat_vmask, ref_14) + assert np.allclose(dreg_patch1_4_lon_vmask, ref_15) + assert np.allclose(dreg_patch1_4_lat_vmask, ref_16) + assert np.allclose(dreg_patch2_1_lon_vmask, ref_17) + assert np.allclose(dreg_patch2_1_lat_vmask, ref_18) + assert np.allclose(dreg_patch2_2_lon_vmask, ref_19) + assert np.allclose(dreg_patch2_2_lat_vmask, ref_20) + assert np.allclose(dreg_patch2_3_lon_vmask, ref_21) + assert np.allclose(dreg_patch2_3_lat_vmask, ref_22) + assert np.allclose(dreg_patch2_4_lon_vmask, ref_23) + assert np.allclose(dreg_patch2_4_lat_vmask, ref_24) diff --git a/advection/tests/test_divide_flux_area_list_stencil_02.py b/advection/tests/test_divide_flux_area_list_stencil_02.py new file mode 100644 index 0000000000..414963e5c0 --- /dev/null +++ b/advection/tests/test_divide_flux_area_list_stencil_02.py @@ -0,0 +1,333 @@ +# 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 + +import numpy as np + +from icon4py.advection.divide_flux_area_list_stencil_02 import ( + divide_flux_area_list_stencil_02, +) +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim + +from .test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from .test_utils.simple_mesh import SimpleMesh + + +def divide_flux_area_list_stencil_02_numpy( + e2c: np.array, + famask_int: np.array, + p_vn: np.array, + bf_cc_patch1_lon: np.array, + bf_cc_patch1_lat: np.array, + bf_cc_patch2_lon: np.array, + bf_cc_patch2_lat: np.array, + butterfly_idx_patch1_vnpos: np.array, + butterfly_idx_patch1_vnneg: np.array, + butterfly_blk_patch1_vnpos: np.array, + butterfly_blk_patch1_vnneg: np.array, + butterfly_idx_patch2_vnpos: np.array, + butterfly_idx_patch2_vnneg: np.array, + butterfly_blk_patch2_vnpos: np.array, + butterfly_blk_patch2_vnneg: np.array, + dreg_patch1_1_lon_vmask: np.array, + dreg_patch1_1_lat_vmask: np.array, + dreg_patch1_2_lon_vmask: np.array, + dreg_patch1_2_lat_vmask: np.array, + dreg_patch1_3_lon_vmask: np.array, + dreg_patch1_3_lat_vmask: np.array, + dreg_patch1_4_lon_vmask: np.array, + dreg_patch1_4_lat_vmask: np.array, + dreg_patch2_1_lon_vmask: np.array, + dreg_patch2_1_lat_vmask: np.array, + dreg_patch2_2_lon_vmask: np.array, + dreg_patch2_2_lat_vmask: np.array, + dreg_patch2_3_lon_vmask: np.array, + dreg_patch2_3_lat_vmask: np.array, + dreg_patch2_4_lon_vmask: np.array, + dreg_patch2_4_lat_vmask: np.array, +): + famask_bool = np.where(famask_int == int32(1), True, False) + lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) + bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) + bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) + bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) + + bf_cc_patch1_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch1_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), + 0.0, + ) + + # patch1 in translated system + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat + + + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape + ) + patch1_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + ) + + +def test_divide_flux_area_list_stencil_02(): + mesh = SimpleMesh() + + famask_int = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_vn = random_field(mesh, EdgeDim, KDim) + bf_cc_patch1_lon = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) + bf_cc_patch1_lat = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) + bf_cc_patch2_lon = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) + bf_cc_patch2_lat = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) + butterfly_idx_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + dreg_patch1_1_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = random_field(mesh, EdgeDim, KDim) + patch1_cell_idx_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch1_cell_blk_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch2_cell_idx_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch2_cell_blk_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20) = divide_flux_area_list_stencil_02_numpy( + mesh.e2c, + np.asarray(famask_int), + np.asarray(p_vn), + np.asarray(bf_cc_patch1_lon), + np.asarray(bf_cc_patch1_lat), + np.asarray(bf_cc_patch2_lon), + np.asarray(bf_cc_patch2_lat), + np.asarray(butterfly_idx_patch1_vnpos), + np.asarray(butterfly_idx_patch1_vnneg), + np.asarray(butterfly_blk_patch1_vnpos), + np.asarray(butterfly_blk_patch1_vnneg), + np.asarray(butterfly_idx_patch2_vnpos), + np.asarray(butterfly_idx_patch2_vnneg), + np.asarray(butterfly_blk_patch2_vnpos), + np.asarray(butterfly_blk_patch2_vnneg), + np.asarray(dreg_patch1_1_lon_vmask), + np.asarray(dreg_patch1_1_lat_vmask), + np.asarray(dreg_patch1_2_lon_vmask), + np.asarray(dreg_patch1_2_lat_vmask), + np.asarray(dreg_patch1_3_lon_vmask), + np.asarray(dreg_patch1_3_lat_vmask), + np.asarray(dreg_patch1_4_lon_vmask), + np.asarray(dreg_patch1_4_lat_vmask), + np.asarray(dreg_patch2_1_lon_vmask), + np.asarray(dreg_patch2_1_lat_vmask), + np.asarray(dreg_patch2_2_lon_vmask), + np.asarray(dreg_patch2_2_lat_vmask), + np.asarray(dreg_patch2_3_lon_vmask), + np.asarray(dreg_patch2_3_lat_vmask), + np.asarray(dreg_patch2_4_lon_vmask), + np.asarray(dreg_patch2_4_lat_vmask), + ) + + divide_flux_area_list_stencil_02( + famask_int, + p_vn, + bf_cc_patch1_lon_field, + bf_cc_patch1_lat_field, + bf_cc_patch2_lon_field, + bf_cc_patch2_lat_field, + butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(dreg_patch1_1_lon_vmask, ref_1) + assert np.allclose(dreg_patch1_1_lat_vmask, ref_2) + assert np.allclose(dreg_patch1_2_lon_vmask, ref_3) + assert np.allclose(dreg_patch1_2_lat_vmask, ref_4) + assert np.allclose(dreg_patch1_3_lon_vmask, ref_5) + assert np.allclose(dreg_patch1_3_lat_vmask, ref_6) + assert np.allclose(dreg_patch1_4_lon_vmask, ref_7) + assert np.allclose(dreg_patch1_4_lat_vmask, ref_8) + assert np.allclose(dreg_patch2_1_lon_vmask, ref_9) + assert np.allclose(dreg_patch2_1_lat_vmask, ref_10) + assert np.allclose(dreg_patch2_2_lon_vmask, ref_11) + assert np.allclose(dreg_patch2_2_lat_vmask, ref_12) + assert np.allclose(dreg_patch2_3_lon_vmask, ref_13) + assert np.allclose(dreg_patch2_3_lat_vmask, ref_14) + assert np.allclose(dreg_patch2_4_lon_vmask, ref_15) + assert np.allclose(dreg_patch2_4_lat_vmask, ref_16) + assert np.allclose(patch1_cell_idx_vmask, ref_17) + assert np.allclose(patch1_cell_blk_vmask, ref_18) + assert np.allclose(patch2_cell_idx_vmask, ref_19) + assert np.allclose(patch2_cell_blk_vmask, ref_20) diff --git a/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py new file mode 100644 index 0000000000..088759ba30 --- /dev/null +++ b/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -0,0 +1,156 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.hflux_ffsl_hybrid_stencil_01a import hflux_ffsl_hybrid_stencil_01a +from icon4py.common.dimension import CellDim, KDim, EdgeDim + +from .test_utils.helpers import _shape, random_field, zero_field, constant_field +from .test_utils.simple_mesh import SimpleMesh + + +def hflux_ffsl_hybrid_stencil_01a_numpy( + e2c: np.array, + z_lsq_coeff_1: np.ndarray, + z_lsq_coeff_2: np.ndarray, + z_lsq_coeff_3: np.ndarray, + z_lsq_coeff_4: np.ndarray, + z_lsq_coeff_5: np.ndarray, + z_lsq_coeff_6: np.ndarray, + z_lsq_coeff_7: np.ndarray, + z_lsq_coeff_8: np.ndarray, + z_lsq_coeff_9: np.ndarray, + z_lsq_coeff_10: np.ndarray, + z_quad_vector_sum0_1: np.ndarray, + z_quad_vector_sum0_2: np.ndarray, + z_quad_vector_sum0_3: np.ndarray, + z_quad_vector_sum0_4: np.ndarray, + z_quad_vector_sum0_5: np.ndarray, + z_quad_vector_sum0_6: np.ndarray, + z_quad_vector_sum0_7: np.ndarray, + z_quad_vector_sum0_8: np.ndarray, + z_quad_vector_sum0_9: np.ndarray, + z_quad_vector_sum0_10: np.ndarray, + patch0_cell_rel_idx_dsl: np.ndarray, +): + + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_hybrid_1a = ( + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) * z_quad_vector_sum0_1 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) * z_quad_vector_sum0_2 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) * z_quad_vector_sum0_3 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_4_e2c[:, 1], z_lsq_coeff_4_e2c[:, 0]) * z_quad_vector_sum0_4 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_5_e2c[:, 1], z_lsq_coeff_5_e2c[:, 0]) * z_quad_vector_sum0_5 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_6_e2c[:, 1], z_lsq_coeff_6_e2c[:, 0]) * z_quad_vector_sum0_6 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_7_e2c[:, 1], z_lsq_coeff_7_e2c[:, 0]) * z_quad_vector_sum0_7 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_8_e2c[:, 1], z_lsq_coeff_8_e2c[:, 0]) * z_quad_vector_sum0_8 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_9_e2c[:, 1], z_lsq_coeff_9_e2c[:, 0]) * z_quad_vector_sum0_9 + + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_10_e2c[:, 1], z_lsq_coeff_10_e2c[:, 0]) * z_quad_vector_sum0_10 + ) + + return p_out_e_hybrid_1a + + +def test_hflux_ffsl_hybrid_stencil_01a(): + mesh = SimpleMesh() + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_4 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_5 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_6 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_7 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_8 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_9 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_10 = random_field(mesh, CellDim, KDim) + z_quad_vector_sum0_1 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_2 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_3 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_4 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_5 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_6 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_7 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_8 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_9 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_10 = random_field(mesh, EdgeDim, KDim) + patch0_cell_rel_idx_dsl = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) + p_out_e_hybrid_1a = zero_field(mesh, EdgeDim, KDim) + + ref = hflux_ffsl_hybrid_stencil_01a_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), + np.asarray(z_lsq_coeff_4), + np.asarray(z_lsq_coeff_5), + np.asarray(z_lsq_coeff_6), + np.asarray(z_lsq_coeff_7), + np.asarray(z_lsq_coeff_8), + np.asarray(z_lsq_coeff_9), + np.asarray(z_lsq_coeff_10), + np.asarray(z_quad_vector_sum0_1), + np.asarray(z_quad_vector_sum0_2), + np.asarray(z_quad_vector_sum0_3), + np.asarray(z_quad_vector_sum0_4), + np.asarray(z_quad_vector_sum0_5), + np.asarray(z_quad_vector_sum0_6), + np.asarray(z_quad_vector_sum0_7), + np.asarray(z_quad_vector_sum0_8), + np.asarray(z_quad_vector_sum0_9), + np.asarray(z_quad_vector_sum0_10), + np.asarray(patch0_cell_rel_idx_dsl), + ) + + hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum0_1, + z_quad_vector_sum0_2, + z_quad_vector_sum0_3, + z_quad_vector_sum0_4, + z_quad_vector_sum0_5, + z_quad_vector_sum0_6, + z_quad_vector_sum0_7, + z_quad_vector_sum0_8, + z_quad_vector_sum0_9, + z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl, + p_out_e_hybrid_1a, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + + ) + + assert np.allclose(p_out_e_hybrid_1a, ref) diff --git a/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py new file mode 100644 index 0000000000..065263b361 --- /dev/null +++ b/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -0,0 +1,55 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.hflux_ffsl_hybrid_stencil_02 import hflux_ffsl_hybrid_stencil_02 +from icon4py.common.dimension import KDim, EdgeDim + +from .test_utils.helpers import _shape, random_field, zero_field, constant_field +from .test_utils.simple_mesh import SimpleMesh + + +def hflux_ffsl_hybrid_stencil_02_numpy( + p_out_e_hybrid_2: np.ndarray, + p_mass_flx_e: np.ndarray, + z_dreg_area: np.ndarray, +): + + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area + + return p_out_e_hybrid_2 + + +def test_hflux_ffsl_hybrid_stencil_02(): + mesh = SimpleMesh() + p_out_e_hybrid_2 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + z_dreg_area = random_field(mesh, EdgeDim, KDim) + + ref = hflux_ffsl_hybrid_stencil_02_numpy( + np.asarray(p_out_e_hybrid_2), + np.asarray(p_mass_flx_e), + np.asarray(z_dreg_area), + ) + + hflux_ffsl_hybrid_stencil_02( + p_out_e_hybrid_2, + p_mass_flx_e, + z_dreg_area, + offset_provider={}, + ) + + assert np.allclose(p_out_e_hybrid_2, ref) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/advection/tests/test_hflx_limiter_mo_stencil_01a.py new file mode 100644 index 0000000000..0b5a4c546d --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_01a.py @@ -0,0 +1,71 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.advection.hflx_limiter_mo_stencil_01a import hflx_limiter_mo_stencil_01a +from icon4py.common.dimension import KDim, EdgeDim, CellDim + +from .test_utils.helpers import _shape, random_field, zero_field, constant_field +from .test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_01a_numpy( + e2c: np.array, + p_mflx_tracer_h: np.ndarray, + p_mass_flx_e: np.ndarray, + p_cc: np.ndarray, +): + p_cc_e2c = p_cc[e2c] + + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) + - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) + ) + + z_anti = p_mflx_tracer_h - z_mflx_low + + return (z_mflx_low, z_anti) + + +def test_hflx_limiter_mo_stencil_01a(): + mesh = SimpleMesh() + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + z_mflx_low = zero_field(mesh, EdgeDim, KDim) + z_anti = zero_field(mesh, EdgeDim, KDim) + + + ref_1, ref_2 = hflx_limiter_mo_stencil_01a_numpy( + mesh.e2c, + np.asarray(p_mflx_tracer_h), + np.asarray(p_mass_flx_e), + np.asarray(p_cc), + ) + + hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h, + p_mass_flx_e, + p_cc, + z_mflx_low, + z_anti, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + + assert np.allclose(z_mflx_low, ref_1) + assert np.allclose(z_anti, ref_2) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/advection/tests/test_hflx_limiter_mo_stencil_01b.py new file mode 100644 index 0000000000..287aa773d9 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_01b.py @@ -0,0 +1,121 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.hflx_limiter_mo_stencil_01b import hflx_limiter_mo_stencil_01b +from icon4py.common.dimension import C2EDim, CEDim, KDim, EdgeDim, CellDim + +from .test_utils.helpers import _shape, random_field, zero_field, constant_field, as_1D_sparse_field +from .test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_01b_numpy( + c2e: np.array, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, + p_dtime: float, +): + z_anti_c2e = z_anti[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + + zero_array = np.zeros(p_rhodz_now.shape) + + z_mflx_anti_1 = ( p_dtime * geofac_div[:, 0] / p_rhodz_new + * z_anti_c2e[:, 0] ) + z_mflx_anti_2 = ( p_dtime * geofac_div[:, 1] / p_rhodz_new + * z_anti_c2e[:, 1] ) + z_mflx_anti_3 = ( p_dtime * geofac_div[:, 2] / p_rhodz_new + * z_anti_c2e[:, 2] ) + + z_mflx_anti_in = -1.0 * ( np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) ) + + z_mflx_anti_out = ( np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) ) + + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) + + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) + + return ( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ) + +def test_hflx_limiter_mo_stencil_01b(): + mesh = SimpleMesh() + + geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_rhodz_now = random_field(mesh, CellDim, KDim) + p_rhodz_new = random_field(mesh, CellDim, KDim) + z_mflx_low = random_field(mesh, EdgeDim, KDim) + z_anti = random_field(mesh, EdgeDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + p_dtime = 5.0 + z_mflx_anti_in = random_field(mesh, CellDim, KDim) + z_mflx_anti_out = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + + ref_1, ref_2, ref_3, ref_4, ref_5 = hflx_limiter_mo_stencil_01b_numpy( + mesh.c2e, + np.asarray(geofac_div), + np.asarray(p_rhodz_now), + np.asarray(p_rhodz_new), + np.asarray(z_mflx_low), + np.asarray(z_anti), + np.asarray(p_cc), + np.asarray(p_dtime) + ) + + hflx_limiter_mo_stencil_01b( + geofac_div_new, + p_rhodz_now, + p_rhodz_new, + z_mflx_low, + z_anti, + p_cc, + p_dtime, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + offset_provider={ + "C2E": mesh.get_c2e_offset_provider(), + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + }, + ) + + assert np.allclose(z_mflx_anti_in, ref_1) + assert np.allclose(z_mflx_anti_out, ref_2) + assert np.allclose(z_tracer_new_low, ref_3) + assert np.allclose(z_tracer_max, ref_4) + assert np.allclose(z_tracer_min, ref_5) diff --git a/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py new file mode 100644 index 0000000000..4783ecaf59 --- /dev/null +++ b/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -0,0 +1,504 @@ +# 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 + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.advection.prep_gauss_quadrature_c_list_stencil import ( + prep_gauss_quadrature_c_list_stencil, +) +from icon4py.common.dimension import EdgeDim, CellDim, KDim + +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field, constant_field +from .test_utils.simple_mesh import SimpleMesh + + +def prep_gauss_quadrature_c_list_stencil_numpy( + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: np.array, +) -> tuple[np.ndarray]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + famask_bool = np.where(famask_int == int32(1), True, False) + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) + + wgt_t_detjac_1 =np.where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_2 =np.where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_3 =np.where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 =np.where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = ( + wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + ) + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**3 + + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + ) + + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + ) + +def test_prep_gauss_quadrature_c_list_stencil(): + mesh = SimpleMesh() + + famask_int = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_dreg_area_in = random_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) + p_dreg_area = zero_field(mesh, EdgeDim, KDim) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11) = prep_gauss_quadrature_c_list_stencil_numpy( + np.asarray(famask_int), + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + np.asarray(p_dreg_area_in), + ) + + prep_gauss_quadrature_c_list_stencil( + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_dreg_area_in, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + offset_provider={}, + ) + co1 = np.asarray(p_quad_vector_sum_1) + co2 = np.asarray(p_quad_vector_sum_2) + co3 = np.asarray(p_quad_vector_sum_3) + co4 = np.asarray(p_quad_vector_sum_4) + co5 = np.asarray(p_quad_vector_sum_5) + co6 = np.asarray(p_quad_vector_sum_6) + co7 = np.asarray(p_quad_vector_sum_7) + co8 = np.asarray(p_quad_vector_sum_8) + co9 = np.asarray(p_quad_vector_sum_9) + co10 = np.asarray(p_quad_vector_sum_10) + co11 = np.asarray(p_dreg_area) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) + assert np.allclose(ref_11, co11) diff --git a/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/advection/tests/test_prep_gauss_quadrature_c_stencil.py new file mode 100644 index 0000000000..9fa2eb5e43 --- /dev/null +++ b/advection/tests/test_prep_gauss_quadrature_c_stencil.py @@ -0,0 +1,464 @@ +# 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 + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.prep_gauss_quadrature_c_stencil import ( + prep_gauss_quadrature_c_stencil, +) +from icon4py.common.dimension import EdgeDim, CellDim, KDim + +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + + +def prep_gauss_quadrature_c_stencil_numpy( + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, +) -> tuple[np.ndarray]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = ( + wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + ) + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**3 + + wgt_t_detjac_2 * z_gauss_pts_2_x**3 + + wgt_t_detjac_3 * z_gauss_pts_3_x**3 + + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y**3 + + wgt_t_detjac_2 * z_gauss_pts_2_y**3 + + wgt_t_detjac_3 * z_gauss_pts_3_y**3 + + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + ) + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = np.where( + z_area >= 0.0, np.maximum(eps, np.absolute(z_area)), -np.maximum(eps, np.absolute(z_area)) + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ) + +def test_prep_gauss_quadrature_c_stencil(): + mesh = SimpleMesh() + + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) + p_dreg_area_out = zero_field(mesh, EdgeDim, KDim) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11) = prep_gauss_quadrature_c_stencil_numpy( + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + ) + + prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + offset_provider={}, + ) + co1 = np.asarray(p_quad_vector_sum_1) + co2 = np.asarray(p_quad_vector_sum_2) + co3 = np.asarray(p_quad_vector_sum_3) + co4 = np.asarray(p_quad_vector_sum_4) + co5 = np.asarray(p_quad_vector_sum_5) + co6 = np.asarray(p_quad_vector_sum_6) + co7 = np.asarray(p_quad_vector_sum_7) + co8 = np.asarray(p_quad_vector_sum_8) + co9 = np.asarray(p_quad_vector_sum_9) + co10 = np.asarray(p_quad_vector_sum_10) + co11 = np.asarray(p_dreg_area_out) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) + assert np.allclose(ref_11, co11) diff --git a/advection/tests/test_recon_lsq_cell_c_stencil.py b/advection/tests/test_recon_lsq_cell_c_stencil.py new file mode 100644 index 0000000000..2e45d69ead --- /dev/null +++ b/advection/tests/test_recon_lsq_cell_c_stencil.py @@ -0,0 +1,683 @@ +# 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 + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.recon_lsq_cell_c_stencil import ( + recon_lsq_cell_c_stencil, +) +from icon4py.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim + +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + + +def recon_lsq_cell_c_stencil_numpy( + c2e2c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_qtmat_c_1: np.ndarray, + lsq_qtmat_c_2: np.ndarray, + lsq_qtmat_c_3: np.ndarray, + lsq_qtmat_c_4: np.ndarray, + lsq_qtmat_c_5: np.ndarray, + lsq_qtmat_c_6: np.ndarray, + lsq_qtmat_c_7: np.ndarray, + lsq_qtmat_c_8: np.ndarray, + lsq_qtmat_c_9: np.ndarray, + lsq_rmat_rdiag_c_1: np.ndarray, + lsq_rmat_rdiag_c_2: np.ndarray, + lsq_rmat_rdiag_c_3: np.ndarray, + lsq_rmat_rdiag_c_4: np.ndarray, + lsq_rmat_rdiag_c_5: np.ndarray, + lsq_rmat_rdiag_c_6: np.ndarray, + lsq_rmat_rdiag_c_7: np.ndarray, + lsq_rmat_rdiag_c_8: np.ndarray, + lsq_rmat_rdiag_c_9: np.ndarray, + lsq_rmat_utri_c_1: np.ndarray, + lsq_rmat_utri_c_2: np.ndarray, + lsq_rmat_utri_c_3: np.ndarray, + lsq_rmat_utri_c_4: np.ndarray, + lsq_rmat_utri_c_5: np.ndarray, + lsq_rmat_utri_c_6: np.ndarray, + lsq_rmat_utri_c_7: np.ndarray, + lsq_rmat_utri_c_8: np.ndarray, + lsq_rmat_utri_c_9: np.ndarray, + lsq_rmat_utri_c_10: np.ndarray, + lsq_rmat_utri_c_11: np.ndarray, + lsq_rmat_utri_c_12: np.ndarray, + lsq_rmat_utri_c_13: np.ndarray, + lsq_rmat_utri_c_14: np.ndarray, + lsq_rmat_utri_c_15: np.ndarray, + lsq_rmat_utri_c_16: np.ndarray, + lsq_rmat_utri_c_17: np.ndarray, + lsq_rmat_utri_c_18: np.ndarray, + lsq_rmat_utri_c_19: np.ndarray, + lsq_rmat_utri_c_20: np.ndarray, + lsq_rmat_utri_c_21: np.ndarray, + lsq_rmat_utri_c_22: np.ndarray, + lsq_rmat_utri_c_23: np.ndarray, + lsq_rmat_utri_c_24: np.ndarray, + lsq_rmat_utri_c_25: np.ndarray, + lsq_rmat_utri_c_26: np.ndarray, + lsq_rmat_utri_c_27: np.ndarray, + lsq_rmat_utri_c_28: np.ndarray, + lsq_rmat_utri_c_29: np.ndarray, + lsq_rmat_utri_c_30: np.ndarray, + lsq_rmat_utri_c_31: np.ndarray, + lsq_rmat_utri_c_32: np.ndarray, + lsq_rmat_utri_c_33: np.ndarray, + lsq_rmat_utri_c_34: np.ndarray, + lsq_rmat_utri_c_35: np.ndarray, + lsq_rmat_utri_c_36: np.ndarray, + lsq_moments_1: np.ndarray, + lsq_moments_2: np.ndarray, + lsq_moments_3: np.ndarray, + lsq_moments_4: np.ndarray, + lsq_moments_5: np.ndarray, + lsq_moments_6: np.ndarray, + lsq_moments_7: np.ndarray, + lsq_moments_8: np.ndarray, + lsq_moments_9: np.ndarray +) -> tuple[np.ndarray]: + p_cc_e = np.expand_dims(p_cc, axis=-1) +# n_diff = p_cc[c2e2c] - p_cc_e + + lsq_rmat_rdiag_c_1 = np.expand_dims(lsq_rmat_rdiag_c_1, axis=-1) + lsq_rmat_rdiag_c_2 = np.expand_dims(lsq_rmat_rdiag_c_2, axis=-1) + lsq_rmat_rdiag_c_3 = np.expand_dims(lsq_rmat_rdiag_c_3, axis=-1) + lsq_rmat_rdiag_c_4 = np.expand_dims(lsq_rmat_rdiag_c_4, axis=-1) + lsq_rmat_rdiag_c_5 = np.expand_dims(lsq_rmat_rdiag_c_5, axis=-1) + lsq_rmat_rdiag_c_6 = np.expand_dims(lsq_rmat_rdiag_c_6, axis=-1) + lsq_rmat_rdiag_c_7 = np.expand_dims(lsq_rmat_rdiag_c_7, axis=-1) + lsq_rmat_rdiag_c_8 = np.expand_dims(lsq_rmat_rdiag_c_8, axis=-1) + lsq_rmat_rdiag_c_9 = np.expand_dims(lsq_rmat_rdiag_c_9, axis=-1) +# lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) +# lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) +# lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) +# lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) +# lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) +# lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) +# lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) +# lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) +# lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) + lsq_rmat_utri_c_1 = np.expand_dims(lsq_rmat_utri_c_1, axis=-1) + lsq_rmat_utri_c_2 = np.expand_dims(lsq_rmat_utri_c_2, axis=-1) + lsq_rmat_utri_c_3 = np.expand_dims(lsq_rmat_utri_c_3, axis=-1) + lsq_rmat_utri_c_4 = np.expand_dims(lsq_rmat_utri_c_4, axis=-1) + lsq_rmat_utri_c_5 = np.expand_dims(lsq_rmat_utri_c_5, axis=-1) + lsq_rmat_utri_c_6 = np.expand_dims(lsq_rmat_utri_c_6, axis=-1) + lsq_rmat_utri_c_7 = np.expand_dims(lsq_rmat_utri_c_7, axis=-1) + lsq_rmat_utri_c_8 = np.expand_dims(lsq_rmat_utri_c_8, axis=-1) + lsq_rmat_utri_c_9 = np.expand_dims(lsq_rmat_utri_c_9, axis=-1) + lsq_rmat_utri_c_10 = np.expand_dims(lsq_rmat_utri_c_10, axis=-1) + lsq_rmat_utri_c_11 = np.expand_dims(lsq_rmat_utri_c_11, axis=-1) + lsq_rmat_utri_c_12 = np.expand_dims(lsq_rmat_utri_c_12, axis=-1) + lsq_rmat_utri_c_13 = np.expand_dims(lsq_rmat_utri_c_13, axis=-1) + lsq_rmat_utri_c_14 = np.expand_dims(lsq_rmat_utri_c_14, axis=-1) + lsq_rmat_utri_c_15 = np.expand_dims(lsq_rmat_utri_c_15, axis=-1) + lsq_rmat_utri_c_16 = np.expand_dims(lsq_rmat_utri_c_16, axis=-1) + lsq_rmat_utri_c_17 = np.expand_dims(lsq_rmat_utri_c_17, axis=-1) + lsq_rmat_utri_c_18 = np.expand_dims(lsq_rmat_utri_c_18, axis=-1) + lsq_rmat_utri_c_19 = np.expand_dims(lsq_rmat_utri_c_19, axis=-1) + lsq_rmat_utri_c_20 = np.expand_dims(lsq_rmat_utri_c_20, axis=-1) + lsq_rmat_utri_c_21 = np.expand_dims(lsq_rmat_utri_c_21, axis=-1) + lsq_rmat_utri_c_22 = np.expand_dims(lsq_rmat_utri_c_22, axis=-1) + lsq_rmat_utri_c_23 = np.expand_dims(lsq_rmat_utri_c_23, axis=-1) + lsq_rmat_utri_c_24 = np.expand_dims(lsq_rmat_utri_c_24, axis=-1) + lsq_rmat_utri_c_25 = np.expand_dims(lsq_rmat_utri_c_25, axis=-1) + lsq_rmat_utri_c_26 = np.expand_dims(lsq_rmat_utri_c_26, axis=-1) + lsq_rmat_utri_c_27 = np.expand_dims(lsq_rmat_utri_c_27, axis=-1) + lsq_rmat_utri_c_28 = np.expand_dims(lsq_rmat_utri_c_28, axis=-1) + lsq_rmat_utri_c_29 = np.expand_dims(lsq_rmat_utri_c_29, axis=-1) + lsq_rmat_utri_c_30 = np.expand_dims(lsq_rmat_utri_c_30, axis=-1) + lsq_rmat_utri_c_31 = np.expand_dims(lsq_rmat_utri_c_31, axis=-1) + lsq_rmat_utri_c_32 = np.expand_dims(lsq_rmat_utri_c_32, axis=-1) + lsq_rmat_utri_c_33 = np.expand_dims(lsq_rmat_utri_c_33, axis=-1) + lsq_rmat_utri_c_34 = np.expand_dims(lsq_rmat_utri_c_34, axis=-1) + lsq_rmat_utri_c_35 = np.expand_dims(lsq_rmat_utri_c_35, axis=-1) + lsq_rmat_utri_c_36 = np.expand_dims(lsq_rmat_utri_c_36, axis=-1) + lsq_rmat_utri_c_1 = np.broadcast_to(lsq_rmat_utri_c_1, p_cc.shape) + lsq_rmat_utri_c_2 = np.broadcast_to(lsq_rmat_utri_c_2, p_cc.shape) + lsq_rmat_utri_c_3 = np.broadcast_to(lsq_rmat_utri_c_3, p_cc.shape) + lsq_rmat_utri_c_4 = np.broadcast_to(lsq_rmat_utri_c_4, p_cc.shape) + lsq_rmat_utri_c_5 = np.broadcast_to(lsq_rmat_utri_c_5, p_cc.shape) + lsq_rmat_utri_c_6 = np.broadcast_to(lsq_rmat_utri_c_6, p_cc.shape) + lsq_rmat_utri_c_7 = np.broadcast_to(lsq_rmat_utri_c_7, p_cc.shape) + lsq_rmat_utri_c_8 = np.broadcast_to(lsq_rmat_utri_c_8, p_cc.shape) + lsq_rmat_utri_c_9 = np.broadcast_to(lsq_rmat_utri_c_9, p_cc.shape) + lsq_rmat_utri_c_10 = np.broadcast_to(lsq_rmat_utri_c_10, p_cc.shape) + lsq_rmat_utri_c_11 = np.broadcast_to(lsq_rmat_utri_c_11, p_cc.shape) + lsq_rmat_utri_c_12 = np.broadcast_to(lsq_rmat_utri_c_12, p_cc.shape) + lsq_rmat_utri_c_13 = np.broadcast_to(lsq_rmat_utri_c_13, p_cc.shape) + lsq_rmat_utri_c_14 = np.broadcast_to(lsq_rmat_utri_c_14, p_cc.shape) + lsq_rmat_utri_c_15 = np.broadcast_to(lsq_rmat_utri_c_15, p_cc.shape) + lsq_rmat_utri_c_16 = np.broadcast_to(lsq_rmat_utri_c_16, p_cc.shape) + lsq_rmat_utri_c_17 = np.broadcast_to(lsq_rmat_utri_c_17, p_cc.shape) + lsq_rmat_utri_c_18 = np.broadcast_to(lsq_rmat_utri_c_18, p_cc.shape) + lsq_rmat_utri_c_19 = np.broadcast_to(lsq_rmat_utri_c_19, p_cc.shape) + lsq_rmat_utri_c_20 = np.broadcast_to(lsq_rmat_utri_c_20, p_cc.shape) + lsq_rmat_utri_c_21 = np.broadcast_to(lsq_rmat_utri_c_21, p_cc.shape) + lsq_rmat_utri_c_22 = np.broadcast_to(lsq_rmat_utri_c_22, p_cc.shape) + lsq_rmat_utri_c_23 = np.broadcast_to(lsq_rmat_utri_c_23, p_cc.shape) + lsq_rmat_utri_c_24 = np.broadcast_to(lsq_rmat_utri_c_24, p_cc.shape) + lsq_rmat_utri_c_25 = np.broadcast_to(lsq_rmat_utri_c_25, p_cc.shape) + lsq_rmat_utri_c_26 = np.broadcast_to(lsq_rmat_utri_c_26, p_cc.shape) + lsq_rmat_utri_c_27 = np.broadcast_to(lsq_rmat_utri_c_27, p_cc.shape) + lsq_rmat_utri_c_28 = np.broadcast_to(lsq_rmat_utri_c_28, p_cc.shape) + lsq_rmat_utri_c_29 = np.broadcast_to(lsq_rmat_utri_c_29, p_cc.shape) + lsq_rmat_utri_c_30 = np.broadcast_to(lsq_rmat_utri_c_30, p_cc.shape) + lsq_rmat_utri_c_31 = np.broadcast_to(lsq_rmat_utri_c_31, p_cc.shape) + lsq_rmat_utri_c_32 = np.broadcast_to(lsq_rmat_utri_c_32, p_cc.shape) + lsq_rmat_utri_c_33 = np.broadcast_to(lsq_rmat_utri_c_33, p_cc.shape) + lsq_rmat_utri_c_34 = np.broadcast_to(lsq_rmat_utri_c_34, p_cc.shape) + lsq_rmat_utri_c_35 = np.broadcast_to(lsq_rmat_utri_c_35, p_cc.shape) + lsq_rmat_utri_c_36 = np.broadcast_to(lsq_rmat_utri_c_36, p_cc.shape) +# lsq_rmat_utri_c_9 lsq_qtmat_c_1 = nplsq_rmat_utri_c_9.broadcast_to(lsq_qtmat_c_1, (CECECDim, KDim)) +# lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, (CECECDim, KDim)) +# lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, (CECECDim, KDim)) +# lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, (CECECDim, KDim)) +# lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, (CECECDim, KDim)) +# lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, (CECECDim, KDim)) +# lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, (CECECDim, KDim)) +# lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, (CECECDim, KDim)) +# lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, (CECECDim, KDim)) +# lsq_qtmat_c_1 = np.broadcast_to(lsq_qtmat_c_1, p_cc.shape) +# lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, p_cc.shape) +# lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, p_cc.shape) +# lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, p_cc.shape) +# lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, p_cc.shape) +# lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, p_cc.shape) +# lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, p_cc.shape) +# lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, p_cc.shape) +# lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, p_cc.shape) +# lsq_qtmat_c_1 = np.repeat(lsq_qtmat_c_1[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_2 = np.repeat(lsq_qtmat_c_2[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_3 = np.repeat(lsq_qtmat_c_3[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_4 = np.repeat(lsq_qtmat_c_4[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_5 = np.repeat(lsq_qtmat_c_5[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_6 = np.repeat(lsq_qtmat_c_6[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_7 = np.repeat(lsq_qtmat_c_7[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_8 = np.repeat(lsq_qtmat_c_8[:, np.newaxis], p_cc.shape[1], axis=1) +# lsq_qtmat_c_9 = np.repeat(lsq_qtmat_c_9[:, np.newaxis], p_cc.shape[1], axis=1) + lsq_qtmat_c_9 = np.expand_dims(lsq_qtmat_c_9, axis=-1) + lsq_qtmat_c_8 = np.expand_dims(lsq_qtmat_c_8, axis=-1) + lsq_qtmat_c_7 = np.expand_dims(lsq_qtmat_c_7, axis=-1) + lsq_qtmat_c_6 = np.expand_dims(lsq_qtmat_c_6, axis=-1) + lsq_qtmat_c_5 = np.expand_dims(lsq_qtmat_c_5, axis=-1) + lsq_qtmat_c_4 = np.expand_dims(lsq_qtmat_c_4, axis=-1) + lsq_qtmat_c_3 = np.expand_dims(lsq_qtmat_c_3, axis=-1) + lsq_qtmat_c_2 = np.expand_dims(lsq_qtmat_c_2, axis=-1) + lsq_qtmat_c_1 = np.expand_dims(lsq_qtmat_c_1, axis=-1) + + p_coeff_10 = lsq_rmat_rdiag_c_9 * ( + lsq_qtmat_c_9[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_9[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_9[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_9[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_9[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_9[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_9[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_9[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_9[:, 8] * (p_cc_e[:, 8] - p_cc) + ) + + + p_coeff_9 = lsq_rmat_rdiag_c_8 * ( + lsq_qtmat_c_8[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_8[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_8[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_8[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_8[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_8[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_8[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_8[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_8[:, 8] * (p_cc_e[:, 8] - p_cc) + - lsq_rmat_utri_c_1 * p_coeff_10 + ) + + p_coeff_8 = lsq_rmat_rdiag_c_8 * ( + lsq_qtmat_c_7[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_7[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_7[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_7[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_7[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_7[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_7[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_7[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_7[:, 8] * (p_cc_e[:, 8] - p_cc) + - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) + ) + + p_coeff_7 = lsq_rmat_rdiag_c_6 * ( + lsq_qtmat_c_6[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_6[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_6[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_6[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_6[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_6[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_6[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_6[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_6[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_4 * p_coeff_8 + + lsq_rmat_utri_c_5 * p_coeff_9 + + lsq_rmat_utri_c_6 * p_coeff_10 + ) + ) + + p_coeff_6 = lsq_rmat_rdiag_c_5 * ( + lsq_qtmat_c_5[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_5[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_5[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_5[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_5[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_5[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_5[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_5[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_5[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_7 * p_coeff_7 + + lsq_rmat_utri_c_8 * p_coeff_8 + + lsq_rmat_utri_c_9 * p_coeff_9 + + lsq_rmat_utri_c_10 * p_coeff_10 + ) + ) + + p_coeff_5 = lsq_rmat_rdiag_c_4 * ( + lsq_qtmat_c_4[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_4[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_4[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_4[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_4[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_4[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_4[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_4[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_4[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_11 * p_coeff_6 + + lsq_rmat_utri_c_12 * p_coeff_7 + + lsq_rmat_utri_c_13 * p_coeff_8 + + lsq_rmat_utri_c_14 * p_coeff_9 + + lsq_rmat_utri_c_15 * p_coeff_10 + ) + ) + + p_coeff_4 = lsq_rmat_rdiag_c_3 * ( + lsq_qtmat_c_3[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_3[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_3[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_3[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_3[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_3[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_3[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_3[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_3[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_16 * p_coeff_5 + + lsq_rmat_utri_c_17 * p_coeff_6 + + lsq_rmat_utri_c_18 * p_coeff_7 + + lsq_rmat_utri_c_19 * p_coeff_8 + + lsq_rmat_utri_c_20 * p_coeff_9 + + lsq_rmat_utri_c_21 * p_coeff_10 + ) + ) + + p_coeff_3 = lsq_rmat_rdiag_c_2 * ( + lsq_qtmat_c_2[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_2[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_2[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_2[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_2[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_2[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_2[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_2[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_2[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_22 * p_coeff_4 + + lsq_rmat_utri_c_23 * p_coeff_5 + + lsq_rmat_utri_c_24 * p_coeff_6 + + lsq_rmat_utri_c_25 * p_coeff_7 + + lsq_rmat_utri_c_26 * p_coeff_8 + + lsq_rmat_utri_c_27 * p_coeff_9 + + lsq_rmat_utri_c_28 * p_coeff_10 + ) + ) + + p_coeff_2 = lsq_rmat_rdiag_c_1 * ( + lsq_qtmat_c_1[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_1[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_1[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_1[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_1[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_1[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_1[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_1[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_1[:, 8] * (p_cc_e[:, 8] - p_cc) + - ( + lsq_rmat_utri_c_29 * p_coeff_3 + + lsq_rmat_utri_c_30 * p_coeff_4 + + lsq_rmat_utri_c_31 * p_coeff_5 + + lsq_rmat_utri_c_32 * p_coeff_6 + + lsq_rmat_utri_c_33 * p_coeff_7 + + lsq_rmat_utri_c_34 * p_coeff_8 + + lsq_rmat_utri_c_35 * p_coeff_9 + + lsq_rmat_utri_c_36 * p_coeff_10 + ) + ) + + p_coeff_1 = p_cc - ( + p_coeff_2 * lsq_moments_1 + + p_coeff_3 * lsq_moments_2 + + p_coeff_4 * lsq_moments_3 + + p_coeff_5 * lsq_moments_4 + + p_coeff_6 * lsq_moments_5 + + p_coeff_7 * lsq_moments_6 + + p_coeff_8 * lsq_moments_7 + + p_coeff_9 * lsq_moments_8 + + p_coeff_10 * lsq_moments_9 + ) + return ( + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + ) + +def test_recon_lsq_cell_c_stencil(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_qtmat_c_1 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_2 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_3 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_4 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_5 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_6 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_7 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_8 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_9 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_qtmat_c_1_field = as_1D_sparse_field(lsq_qtmat_c_1, CECECDim) + lsq_qtmat_c_2_field = as_1D_sparse_field(lsq_qtmat_c_2, CECECDim) + lsq_qtmat_c_3_field = as_1D_sparse_field(lsq_qtmat_c_3, CECECDim) + lsq_qtmat_c_4_field = as_1D_sparse_field(lsq_qtmat_c_4, CECECDim) + lsq_qtmat_c_5_field = as_1D_sparse_field(lsq_qtmat_c_5, CECECDim) + lsq_qtmat_c_6_field = as_1D_sparse_field(lsq_qtmat_c_6, CECECDim) + lsq_qtmat_c_7_field = as_1D_sparse_field(lsq_qtmat_c_7, CECECDim) + lsq_qtmat_c_8_field = as_1D_sparse_field(lsq_qtmat_c_8, CECECDim) + lsq_qtmat_c_9_field = as_1D_sparse_field(lsq_qtmat_c_9, CECECDim) + lsq_rmat_rdiag_c_1 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_2 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_3 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_4 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_5 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_6 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_7 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_8 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_9 = random_field(mesh, CellDim) + lsq_rmat_utri_c_1 = random_field(mesh, CellDim) + lsq_rmat_utri_c_2 = random_field(mesh, CellDim) + lsq_rmat_utri_c_3 = random_field(mesh, CellDim) + lsq_rmat_utri_c_4 = random_field(mesh, CellDim) + lsq_rmat_utri_c_5 = random_field(mesh, CellDim) + lsq_rmat_utri_c_6 = random_field(mesh, CellDim) + lsq_rmat_utri_c_7 = random_field(mesh, CellDim) + lsq_rmat_utri_c_8 = random_field(mesh, CellDim) + lsq_rmat_utri_c_9 = random_field(mesh, CellDim) + lsq_rmat_utri_c_10 = random_field(mesh, CellDim) + lsq_rmat_utri_c_11 = random_field(mesh, CellDim) + lsq_rmat_utri_c_12 = random_field(mesh, CellDim) + lsq_rmat_utri_c_13 = random_field(mesh, CellDim) + lsq_rmat_utri_c_14 = random_field(mesh, CellDim) + lsq_rmat_utri_c_15 = random_field(mesh, CellDim) + lsq_rmat_utri_c_16 = random_field(mesh, CellDim) + lsq_rmat_utri_c_17 = random_field(mesh, CellDim) + lsq_rmat_utri_c_18 = random_field(mesh, CellDim) + lsq_rmat_utri_c_19 = random_field(mesh, CellDim) + lsq_rmat_utri_c_20 = random_field(mesh, CellDim) + lsq_rmat_utri_c_21 = random_field(mesh, CellDim) + lsq_rmat_utri_c_22 = random_field(mesh, CellDim) + lsq_rmat_utri_c_23 = random_field(mesh, CellDim) + lsq_rmat_utri_c_24 = random_field(mesh, CellDim) + lsq_rmat_utri_c_25 = random_field(mesh, CellDim) + lsq_rmat_utri_c_26 = random_field(mesh, CellDim) + lsq_rmat_utri_c_27 = random_field(mesh, CellDim) + lsq_rmat_utri_c_28 = random_field(mesh, CellDim) + lsq_rmat_utri_c_29 = random_field(mesh, CellDim) + lsq_rmat_utri_c_30 = random_field(mesh, CellDim) + lsq_rmat_utri_c_31 = random_field(mesh, CellDim) + lsq_rmat_utri_c_32 = random_field(mesh, CellDim) + lsq_rmat_utri_c_33 = random_field(mesh, CellDim) + lsq_rmat_utri_c_34 = random_field(mesh, CellDim) + lsq_rmat_utri_c_35 = random_field(mesh, CellDim) + lsq_rmat_utri_c_36 = random_field(mesh, CellDim) + lsq_moments_1 = random_field(mesh, CellDim) + lsq_moments_2 = random_field(mesh, CellDim) + lsq_moments_3 = random_field(mesh, CellDim) + lsq_moments_4 = random_field(mesh, CellDim) + lsq_moments_5 = random_field(mesh, CellDim) + lsq_moments_6 = random_field(mesh, CellDim) + lsq_moments_7 = random_field(mesh, CellDim) + lsq_moments_8 = random_field(mesh, CellDim) + lsq_moments_9 = random_field(mesh, CellDim) + p_coeff_1 = zero_field(mesh, CellDim, KDim) + p_coeff_2 = zero_field(mesh, CellDim, KDim) + p_coeff_3 = zero_field(mesh, CellDim, KDim) + p_coeff_4 = zero_field(mesh, CellDim, KDim) + p_coeff_5 = zero_field(mesh, CellDim, KDim) + p_coeff_6 = zero_field(mesh, CellDim, KDim) + p_coeff_7 = zero_field(mesh, CellDim, KDim) + p_coeff_8 = zero_field(mesh, CellDim, KDim) + p_coeff_9 = zero_field(mesh, CellDim, KDim) + p_coeff_10 = zero_field(mesh, CellDim, KDim) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10) = recon_lsq_cell_c_stencil_numpy( + mesh.c2e2c2e2c, + np.asarray(p_cc), + np.asarray(lsq_qtmat_c_1), + np.asarray(lsq_qtmat_c_2), + np.asarray(lsq_qtmat_c_3), + np.asarray(lsq_qtmat_c_4), + np.asarray(lsq_qtmat_c_5), + np.asarray(lsq_qtmat_c_6), + np.asarray(lsq_qtmat_c_7), + np.asarray(lsq_qtmat_c_8), + np.asarray(lsq_qtmat_c_9), + np.asarray(lsq_rmat_rdiag_c_1), + np.asarray(lsq_rmat_rdiag_c_2), + np.asarray(lsq_rmat_rdiag_c_3), + np.asarray(lsq_rmat_rdiag_c_4), + np.asarray(lsq_rmat_rdiag_c_5), + np.asarray(lsq_rmat_rdiag_c_6), + np.asarray(lsq_rmat_rdiag_c_7), + np.asarray(lsq_rmat_rdiag_c_8), + np.asarray(lsq_rmat_rdiag_c_9), + np.asarray(lsq_rmat_utri_c_1), + np.asarray(lsq_rmat_utri_c_2), + np.asarray(lsq_rmat_utri_c_3), + np.asarray(lsq_rmat_utri_c_4), + np.asarray(lsq_rmat_utri_c_5), + np.asarray(lsq_rmat_utri_c_6), + np.asarray(lsq_rmat_utri_c_7), + np.asarray(lsq_rmat_utri_c_8), + np.asarray(lsq_rmat_utri_c_9), + np.asarray(lsq_rmat_utri_c_10), + np.asarray(lsq_rmat_utri_c_11), + np.asarray(lsq_rmat_utri_c_12), + np.asarray(lsq_rmat_utri_c_13), + np.asarray(lsq_rmat_utri_c_14), + np.asarray(lsq_rmat_utri_c_15), + np.asarray(lsq_rmat_utri_c_16), + np.asarray(lsq_rmat_utri_c_17), + np.asarray(lsq_rmat_utri_c_18), + np.asarray(lsq_rmat_utri_c_19), + np.asarray(lsq_rmat_utri_c_20), + np.asarray(lsq_rmat_utri_c_21), + np.asarray(lsq_rmat_utri_c_22), + np.asarray(lsq_rmat_utri_c_23), + np.asarray(lsq_rmat_utri_c_24), + np.asarray(lsq_rmat_utri_c_25), + np.asarray(lsq_rmat_utri_c_26), + np.asarray(lsq_rmat_utri_c_27), + np.asarray(lsq_rmat_utri_c_28), + np.asarray(lsq_rmat_utri_c_29), + np.asarray(lsq_rmat_utri_c_30), + np.asarray(lsq_rmat_utri_c_31), + np.asarray(lsq_rmat_utri_c_32), + np.asarray(lsq_rmat_utri_c_33), + np.asarray(lsq_rmat_utri_c_34), + np.asarray(lsq_rmat_utri_c_35), + np.asarray(lsq_rmat_utri_c_36), + np.asarray(lsq_moments_1), + np.asarray(lsq_moments_2), + np.asarray(lsq_moments_3), + np.asarray(lsq_moments_4), + np.asarray(lsq_moments_5), + np.asarray(lsq_moments_6), + np.asarray(lsq_moments_7), + np.asarray(lsq_moments_8), + np.asarray(lsq_moments_9) + ) + + recon_lsq_cell_c_stencil( + p_cc, + lsq_qtmat_c_1_field, + lsq_qtmat_c_2_field, + lsq_qtmat_c_3_field, + lsq_qtmat_c_4_field, + lsq_qtmat_c_5_field, + lsq_qtmat_c_6_field, + lsq_qtmat_c_7_field, + lsq_qtmat_c_8_field, + lsq_qtmat_c_9_field, + lsq_rmat_rdiag_c_1, + lsq_rmat_rdiag_c_2, + lsq_rmat_rdiag_c_3, + lsq_rmat_rdiag_c_4, + lsq_rmat_rdiag_c_5, + lsq_rmat_rdiag_c_6, + lsq_rmat_rdiag_c_7, + lsq_rmat_rdiag_c_8, + lsq_rmat_rdiag_c_9, + lsq_rmat_utri_c_1, + lsq_rmat_utri_c_2, + lsq_rmat_utri_c_3, + lsq_rmat_utri_c_4, + lsq_rmat_utri_c_5, + lsq_rmat_utri_c_6, + lsq_rmat_utri_c_7, + lsq_rmat_utri_c_8, + lsq_rmat_utri_c_9, + lsq_rmat_utri_c_10, + lsq_rmat_utri_c_11, + lsq_rmat_utri_c_12, + lsq_rmat_utri_c_13, + lsq_rmat_utri_c_14, + lsq_rmat_utri_c_15, + lsq_rmat_utri_c_16, + lsq_rmat_utri_c_17, + lsq_rmat_utri_c_18, + lsq_rmat_utri_c_19, + lsq_rmat_utri_c_20, + lsq_rmat_utri_c_21, + lsq_rmat_utri_c_22, + lsq_rmat_utri_c_23, + lsq_rmat_utri_c_24, + lsq_rmat_utri_c_25, + lsq_rmat_utri_c_26, + lsq_rmat_utri_c_27, + lsq_rmat_utri_c_28, + lsq_rmat_utri_c_29, + lsq_rmat_utri_c_30, + lsq_rmat_utri_c_31, + lsq_rmat_utri_c_32, + lsq_rmat_utri_c_33, + lsq_rmat_utri_c_34, + lsq_rmat_utri_c_35, + lsq_rmat_utri_c_36, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + p_coeff_1, + p_coeff_2, + p_coeff_3, + p_coeff_4, + p_coeff_5, + p_coeff_6, + p_coeff_7, + p_coeff_8, + p_coeff_9, + p_coeff_10, + offset_provider={ + "C2E2C2E2C": mesh.get_c2e2c_offset_provider(), + "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), + }, + ) + co1 = np.asarray(p_coeff_1) + co2 = np.asarray(p_coeff_2) + co3 = np.asarray(p_coeff_3) + co4 = np.asarray(p_coeff_4) + co5 = np.asarray(p_coeff_5) + co6 = np.asarray(p_coeff_6) + co7 = np.asarray(p_coeff_7) + co8 = np.asarray(p_coeff_8) + co9 = np.asarray(p_coeff_9) + co10 = np.asarray(p_coeff_10) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) diff --git a/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/advection/tests/test_recon_lsq_cell_c_svd_stencil.py new file mode 100644 index 0000000000..6246fd4520 --- /dev/null +++ b/advection/tests/test_recon_lsq_cell_c_svd_stencil.py @@ -0,0 +1,337 @@ +# 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 + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.recon_lsq_cell_c_svd_stencil import ( + recon_lsq_cell_c_svd_stencil, +) +from icon4py.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim + +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + +def recon_lsq_cell_c_svd_stencil_numpy( + c2e2c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_pseudoinv_1: np.ndarray, + lsq_pseudoinv_2: np.ndarray, + lsq_pseudoinv_3: np.ndarray, + lsq_pseudoinv_4: np.ndarray, + lsq_pseudoinv_5: np.ndarray, + lsq_pseudoinv_6: np.ndarray, + lsq_pseudoinv_7: np.ndarray, + lsq_pseudoinv_8: np.ndarray, + lsq_pseudoinv_9: np.ndarray, + lsq_moments_1: np.ndarray, + lsq_moments_2: np.ndarray, + lsq_moments_3: np.ndarray, + lsq_moments_4: np.ndarray, + lsq_moments_5: np.ndarray, + lsq_moments_6: np.ndarray, + lsq_moments_7: np.ndarray, + lsq_moments_8: np.ndarray, + lsq_moments_9: np.ndarray +) -> tuple[np.ndarray]: + + p_cc_e = np.expand_dims(p_cc, axis=-1) + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + + p_coeff_10_dsl = ( + lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + + p_coeff_9_dsl = ( + lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_8_dsl = ( + lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_7_dsl = ( + lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_6_dsl = ( + lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_5_dsl = ( + lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_4_dsl = ( + lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_3_dsl = ( + lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_2_dsl = ( + lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 + ) + return ( + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + ) + +def test_recon_lsq_cell_c_svd_stencil(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_2 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_3 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_4 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_5 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_6 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_7 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_8 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_9 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECECDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECECDim) + lsq_pseudoinv_3_field = as_1D_sparse_field(lsq_pseudoinv_3, CECECDim) + lsq_pseudoinv_4_field = as_1D_sparse_field(lsq_pseudoinv_4, CECECDim) + lsq_pseudoinv_5_field = as_1D_sparse_field(lsq_pseudoinv_5, CECECDim) + lsq_pseudoinv_6_field = as_1D_sparse_field(lsq_pseudoinv_6, CECECDim) + lsq_pseudoinv_7_field = as_1D_sparse_field(lsq_pseudoinv_7, CECECDim) + lsq_pseudoinv_8_field = as_1D_sparse_field(lsq_pseudoinv_8, CECECDim) + lsq_pseudoinv_9_field = as_1D_sparse_field(lsq_pseudoinv_9, CECECDim) + lsq_moments_1 = random_field(mesh, CellDim) + lsq_moments_2 = random_field(mesh, CellDim) + lsq_moments_3 = random_field(mesh, CellDim) + lsq_moments_4 = random_field(mesh, CellDim) + lsq_moments_5 = random_field(mesh, CellDim) + lsq_moments_6 = random_field(mesh, CellDim) + lsq_moments_7 = random_field(mesh, CellDim) + lsq_moments_8 = random_field(mesh, CellDim) + lsq_moments_9 = random_field(mesh, CellDim) + p_coeff_1_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_2_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_3_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_4_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_5_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_6_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_7_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_8_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_9_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_10_dsl = zero_field(mesh, CellDim, KDim) + + (ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10) = recon_lsq_cell_c_svd_stencil_numpy( + mesh.c2e2c2e2c, + np.asarray(p_cc), + np.asarray(lsq_pseudoinv_1), + np.asarray(lsq_pseudoinv_2), + np.asarray(lsq_pseudoinv_3), + np.asarray(lsq_pseudoinv_4), + np.asarray(lsq_pseudoinv_5), + np.asarray(lsq_pseudoinv_6), + np.asarray(lsq_pseudoinv_7), + np.asarray(lsq_pseudoinv_8), + np.asarray(lsq_pseudoinv_9), + np.asarray(lsq_moments_1), + np.asarray(lsq_moments_2), + np.asarray(lsq_moments_3), + np.asarray(lsq_moments_4), + np.asarray(lsq_moments_5), + np.asarray(lsq_moments_6), + np.asarray(lsq_moments_7), + np.asarray(lsq_moments_8), + np.asarray(lsq_moments_9) + ) + + recon_lsq_cell_c_svd_stencil( + p_cc, + lsq_pseudoinv_1_field, + lsq_pseudoinv_2_field, + lsq_pseudoinv_3_field, + lsq_pseudoinv_4_field, + lsq_pseudoinv_5_field, + lsq_pseudoinv_6_field, + lsq_pseudoinv_7_field, + lsq_pseudoinv_8_field, + lsq_pseudoinv_9_field, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + offset_provider={ + "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), + "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), + }, + ) +# co1 = np.asarray(p_coeff_1_dsl) + co2 = np.asarray(p_coeff_2_dsl) + co3 = np.asarray(p_coeff_3_dsl) + co4 = np.asarray(p_coeff_4_dsl) + co5 = np.asarray(p_coeff_5_dsl) + co6 = np.asarray(p_coeff_6_dsl) + co7 = np.asarray(p_coeff_7_dsl) + co8 = np.asarray(p_coeff_8_dsl) + co9 = np.asarray(p_coeff_9_dsl) + co10 = np.asarray(p_coeff_10_dsl) +# assert np.allclose(ref_1, co1) + assert np.allclose(ref_10, co10) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_recon_lsq_cell_l_svd_stencil.py similarity index 89% rename from advection/tests/test_upwind_hflux_miura_stencil_02.py rename to advection/tests/test_recon_lsq_cell_l_svd_stencil.py index d748eae380..71b8aeff0f 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_02.py +++ b/advection/tests/test_recon_lsq_cell_l_svd_stencil.py @@ -14,8 +14,8 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.upwind_hflux_miura_stencil_02 import ( - upwind_hflux_miura_stencil_02, +from icon4py.advection.recon_lsq_cell_l_svd_stencil import ( + recon_lsq_cell_l_svd_stencil, ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim @@ -23,7 +23,7 @@ from .test_utils.simple_mesh import SimpleMesh -def upwind_hflux_miura_stencil_02_numpy( +def recon_lsq_cell_l_svd_stencil_numpy( c2e2c: np.ndarray, p_cc: np.ndarray, lsq_pseudoinv_1: np.ndarray, @@ -39,7 +39,7 @@ def upwind_hflux_miura_stencil_02_numpy( return p_coeff_1, p_coeff_2, p_coeff_3 -def test_upwind_hflux_miura_stencil_02(): +def test_recon_lsq_cell_l_svd_stencil(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2CDim) @@ -51,14 +51,14 @@ def test_upwind_hflux_miura_stencil_02(): p_coeff_2 = zero_field(mesh, CellDim, KDim) p_coeff_3 = zero_field(mesh, CellDim, KDim) - ref_1, ref_2, ref_3 = upwind_hflux_miura_stencil_02_numpy( + ref_1, ref_2, ref_3 = recon_lsq_cell_l_svd_stencil_numpy( mesh.c2e2c, np.asarray(p_cc), np.asarray(lsq_pseudoinv_1), np.asarray(lsq_pseudoinv_2), ) - upwind_hflux_miura_stencil_02( + recon_lsq_cell_l_svd_stencil( p_cc, lsq_pseudoinv_1_field, lsq_pseudoinv_2_field, diff --git a/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/advection/tests/test_upwind_hflux_miura3_stencil_01.py new file mode 100644 index 0000000000..f1b1519a51 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura3_stencil_01.py @@ -0,0 +1,165 @@ +# 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 + +import numpy as np + +from icon4py.advection.upwind_hflux_miura3_stencil_01 import ( + upwind_hflux_miura3_stencil_01, +) +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.common.dimension import CellDim, EdgeDim, KDim + +from .test_utils.helpers import random_field, zero_field, random_mask +from .test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura3_stencil_01_numpy( + e2c: np.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum_1: np.array, + z_quad_vector_sum_2: np.array, + z_quad_vector_sum_3: np.array, + z_quad_vector_sum_4: np.array, + z_quad_vector_sum_5: np.array, + z_quad_vector_sum_6: np.array, + z_quad_vector_sum_7: np.array, + z_quad_vector_sum_8: np.array, + z_quad_vector_sum_9: np.array, + z_quad_vector_sum_10: np.array, + z_dreg_area: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, +) -> np.array: + + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_miura3 = ( + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) * z_quad_vector_sum_1 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) * z_quad_vector_sum_2 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) * z_quad_vector_sum_3 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_4_e2c[:, 1], z_lsq_coeff_4_e2c[:, 0]) * z_quad_vector_sum_4 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_5_e2c[:, 1], z_lsq_coeff_5_e2c[:, 0]) * z_quad_vector_sum_5 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_6_e2c[:, 1], z_lsq_coeff_6_e2c[:, 0]) * z_quad_vector_sum_6 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_7_e2c[:, 1], z_lsq_coeff_7_e2c[:, 0]) * z_quad_vector_sum_7 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_8_e2c[:, 1], z_lsq_coeff_8_e2c[:, 0]) * z_quad_vector_sum_8 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_9_e2c[:, 1], z_lsq_coeff_9_e2c[:, 0]) * z_quad_vector_sum_9 + + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_10_e2c[:, 1], z_lsq_coeff_10_e2c[:, 0]) * z_quad_vector_sum_10 + ) / z_dreg_area * p_mass_flx_e + + return p_out_e_miura3 + + +def test_upwind_hflux_miura3_stencil_01(): + mesh = SimpleMesh() + + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_4 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_5 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_6 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_7 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_8 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_9 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_10 = random_field(mesh, CellDim, KDim) + z_quad_vector_sum_1 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_2 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_3 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_4 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_5 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_6 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_7 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_8 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_9 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_10 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + z_dreg_area = random_field(mesh, EdgeDim, KDim) +# cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_out_e_miura3 = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura3_stencil_01_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), + np.asarray(z_lsq_coeff_4), + np.asarray(z_lsq_coeff_5), + np.asarray(z_lsq_coeff_6), + np.asarray(z_lsq_coeff_7), + np.asarray(z_lsq_coeff_8), + np.asarray(z_lsq_coeff_9), + np.asarray(z_lsq_coeff_10), + np.asarray(z_quad_vector_sum_1), + np.asarray(z_quad_vector_sum_2), + np.asarray(z_quad_vector_sum_3), + np.asarray(z_quad_vector_sum_4), + np.asarray(z_quad_vector_sum_5), + np.asarray(z_quad_vector_sum_6), + np.asarray(z_quad_vector_sum_7), + np.asarray(z_quad_vector_sum_8), + np.asarray(z_quad_vector_sum_9), + np.asarray(z_quad_vector_sum_10), + np.asarray(z_dreg_area), + np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), + ) + + upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum_1, + z_quad_vector_sum_2, + z_quad_vector_sum_3, + z_quad_vector_sum_4, + z_quad_vector_sum_5, + z_quad_vector_sum_6, + z_quad_vector_sum_7, + z_quad_vector_sum_8, + z_quad_vector_sum_9, + z_quad_vector_sum_10, + z_dreg_area, + p_mass_flx_e, + cell_rel_idx_dsl, + p_out_e_miura3, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(np.asarray(p_out_e_miura3), ref) diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py new file mode 100644 index 0000000000..59cf65e96b --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -0,0 +1,82 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.advection.upwind_hflux_miura_cycl_stencil_01 import upwind_hflux_miura_cycl_stencil_01 +from icon4py.common.dimension import CellDim, EdgeDim, KDim + +from .test_utils.helpers import _shape, random_field, zero_field, random_mask +from .test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_01_numpy( + e2c: np.array, + z_lsq_coeff_1_dsl: np.array, + z_lsq_coeff_2_dsl: np.array, + z_lsq_coeff_3_dsl: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: int32, +): + z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] + z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] + z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] + + z_tracer_mflx_dsl = ( np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl_e2c[:, 1], z_lsq_coeff_1_dsl_e2c[:, 0]) + + distv_bary_1 * np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_dsl_e2c[:, 1], z_lsq_coeff_2_dsl_e2c[:, 0]) + + distv_bary_2 * np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_dsl_e2c[:, 1], z_lsq_coeff_3_dsl_e2c[:, 0]) + ) * p_mass_flx_e + + return z_tracer_mflx_dsl + + +def test_upwind_hflux_miura_cycl_stencil_01(): + mesh = SimpleMesh() + + z_lsq_coeff_1_dsl = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2_dsl = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3_dsl = random_field(mesh, CellDim, KDim) + distv_bary_1 = random_field(mesh, EdgeDim, KDim) + distv_bary_2 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) + z_tracer_mflx_dsl = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_01_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1_dsl), + np.asarray(z_lsq_coeff_2_dsl), + np.asarray(z_lsq_coeff_3_dsl), + np.asarray(distv_bary_1), + np.asarray(distv_bary_2), + np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), + ) + + upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + z_tracer_mflx_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(ref, z_tracer_mflx_dsl) diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py new file mode 100644 index 0000000000..ca075b3022 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -0,0 +1,102 @@ +# 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 + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.upwind_hflux_miura_cycl_stencil_02 import upwind_hflux_miura_cycl_stencil_02 +from icon4py.common.dimension import CellDim, EdgeDim, KDim, CEDim, C2EDim + +from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from .test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_02_numpy( + c2e: np.array, + nsub: int32, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, + z_dtsub: float, +): + p_mass_flx_e_c2e = p_mass_flx_e[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_tracer_mflx_c2e = z_tracer_mflx[c2e] + + z_rhofluxdiv_c_out = np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + + z_rho_new_dsl = ( z_rho_now + - z_dtsub * z_rhofluxdiv_c_out ) + + z_tracer_new_dsl = (( z_tracer_now + * z_rho_now + - z_dtsub * z_fluxdiv_c_dsl ) + / z_rho_new_dsl) + + return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + +def test_upwind_hflux_miura_cycl_stencil_02(): + mesh = SimpleMesh() + nsub = int32(1) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div_field = as_1D_sparse_field(geofac_div, CEDim) + z_rhofluxdiv_c = random_field(mesh, CellDim, KDim) + z_tracer_mflx = random_field(mesh, EdgeDim, KDim) + z_rho_now = random_field(mesh, CellDim, KDim) + z_tracer_now = random_field(mesh, CellDim, KDim) + z_dtsub = 0.5 + z_rhofluxdiv_c_out = random_field(mesh, CellDim, KDim) + z_fluxdiv_c_dsl = random_field(mesh, CellDim, KDim) + z_rho_new_dsl = random_field(mesh, CellDim, KDim) + z_tracer_new_dsl = random_field(mesh, CellDim, KDim) + + ref_1, ref_2, ref_3, ref_4 = upwind_hflux_miura_cycl_stencil_02_numpy( + mesh.c2e, + nsub, + np.asarray(p_mass_flx_e), + np.asarray(geofac_div), + np.asarray(z_rhofluxdiv_c), + np.asarray(z_tracer_mflx), + np.asarray(z_rho_now), + np.asarray(z_tracer_now), + z_dtsub, + ) + + upwind_hflux_miura_cycl_stencil_02( + nsub, + p_mass_flx_e, + geofac_div, + z_rhofluxdiv_c, + z_tracer_mflx, + z_rho_now, + z_tracer_now, + z_dtsub, + z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl, + z_rho_new_dsl, + z_tracer_new_dsl, + offset_provider={ + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + "C2E": mesh.get_c2e_offset_provider(), + }, + ) + assert np.allclose(ref_1, z_rhofluxdiv_c_out) + assert np.allclose(ref_2, z_fluxdiv_c_dsl) + assert np.allclose(ref_3, z_rho_new_dsl) + assert np.allclose(ref_4, z_tracer_new_dsl) diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py new file mode 100644 index 0000000000..1c3e005af5 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -0,0 +1,47 @@ +# 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 + +import numpy as np + +from icon4py.advection.upwind_hflux_miura_cycl_stencil_03a import upwind_hflux_miura_cycl_stencil_03a +from icon4py.common.dimension import EdgeDim, KDim + +from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_03a_numpy( + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, +): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl)/float(2) + return p_out_e + +def test_upwind_hflux_miura_cycl_stencil_03a(): + mesh = SimpleMesh() + z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(mesh, EdgeDim, KDim) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_03a_numpy( + np.asarray(z_tracer_mflx_1_dsl), + np.asarray(z_tracer_mflx_2_dsl), + ) + + upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + p_out_e, + offset_provider={}, + ) + assert np.allclose(ref, p_out_e) diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py new file mode 100644 index 0000000000..2e8df339ed --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -0,0 +1,53 @@ +# 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 + +import numpy as np + +from icon4py.advection.upwind_hflux_miura_cycl_stencil_03b import upwind_hflux_miura_cycl_stencil_03b +from icon4py.common.dimension import EdgeDim, KDim + +from .test_utils.helpers import random_field, zero_field +from .test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_03b_numpy( + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, +): + p_out_e = (z_tracer_mflx_1_dsl + + z_tracer_mflx_2_dsl + + z_tracer_mflx_3_dsl)/float(3) + return p_out_e + +def test_upwind_hflux_miura_cycl_stencil_03b(): + mesh = SimpleMesh() + z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_3_dsl = random_field(mesh, EdgeDim, KDim) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_03b_numpy( + np.asarray(z_tracer_mflx_1_dsl), + np.asarray(z_tracer_mflx_2_dsl), + np.asarray(z_tracer_mflx_3_dsl), + ) + + upwind_hflux_miura_cycl_stencil_03b( + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + z_tracer_mflx_3_dsl, + p_out_e, + offset_provider={}, + ) + assert np.allclose(ref, p_out_e) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/advection/tests/test_upwind_hflux_miura_stencil_01.py index 868ffc38dc..faf80d6030 100644 --- a/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -16,42 +16,32 @@ from icon4py.advection.upwind_hflux_miura_stencil_01 import ( upwind_hflux_miura_stencil_01, ) +from gt4py.next.ffront.fbuiltins import int32 from icon4py.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field +from .test_utils.helpers import random_field, zero_field, constant_field from .test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_stencil_01_numpy( e2c: np.array, - p_vn: np.array, - p_cc: np.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, distv_bary_1: np.array, distv_bary_2: np.array, - z_grad_1: np.array, - z_grad_2: np.array, p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, ) -> np.array: - p_cc_e2c = p_cc[e2c] - z_grad_1_e2c = z_grad_1[e2c] - z_grad_2_e2c = z_grad_2[e2c] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - p_out_e = np.where( - p_vn > 0.0, - ( - p_cc_e2c[:, 0] - + distv_bary_1 * z_grad_1_e2c[:, 0] - + distv_bary_2 * z_grad_2_e2c[:, 0] - ) - * p_mass_flx_e, - ( - p_cc_e2c[:, 1] - + distv_bary_1 * z_grad_1_e2c[:, 1] - + distv_bary_2 * z_grad_2_e2c[:, 1] - ) - * p_mass_flx_e, - ) + p_out_e = ( np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + + distv_bary_1 *np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + + distv_bary_2 *np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + ) * p_mass_flx_e return p_out_e @@ -59,34 +49,34 @@ def upwind_hflux_miura_stencil_01_numpy( def test_upwind_hflux_miura_stencil_01(): mesh = SimpleMesh() - p_vn = random_field(mesh, EdgeDim, KDim) - p_cc = random_field(mesh, CellDim, KDim) + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) distv_bary_1 = random_field(mesh, EdgeDim, KDim) distv_bary_2 = random_field(mesh, EdgeDim, KDim) - z_grad_1 = random_field(mesh, CellDim, KDim) - z_grad_2 = random_field(mesh, CellDim, KDim) p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_out_e = zero_field(mesh, EdgeDim, KDim) ref = upwind_hflux_miura_stencil_01_numpy( mesh.e2c, - np.asarray(p_vn), - np.asarray(p_cc), + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), np.asarray(distv_bary_1), np.asarray(distv_bary_2), - np.asarray(z_grad_1), - np.asarray(z_grad_2), np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), ) upwind_hflux_miura_stencil_01( - p_vn, - p_cc, + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, distv_bary_1, distv_bary_2, - z_grad_1, - z_grad_2, p_mass_flx_e, + cell_rel_idx_dsl, p_out_e, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), diff --git a/advection/tests/test_utils/simple_mesh.py b/advection/tests/test_utils/simple_mesh.py index 5b66aeed19..1204f288c4 100644 --- a/advection/tests/test_utils/simple_mesh.py +++ b/advection/tests/test_utils/simple_mesh.py @@ -17,6 +17,7 @@ from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from icon4py.common.dimension import ( + C2E2C2E2CDim, C2E2CDim, C2E2CODim, C2EDim, @@ -347,6 +348,28 @@ class SimpleMeshData: ] ) + c2e2c2e2c_table = np.asarray( + [ + [15, 4, 3,12, 14, 1, 7, 6, 2], #1c + [16, 5, 4,12, 13, 2, 8, 7, 0], + [17, 3, 5,13, 14, 0, 6, 8, 1], + [0, 6, 2, 17, 5, 9, 10, 15, 4], + [1, 7, 0, 15, 3, 16, 5, 10, 11], + [2, 8, 1, 4, 16, 17, 3, 9, 11], #5c + [3, 10, 9, 2, 0, 7, 13, 8, 12], + [4, 11, 10, 0, 1, 8, 14, 6, 13], + [5, 9, 11, 1, 2, 3, 12, 7, 14], + [6, 12, 8, 5, 11, 3, 10, 16, 15], + [7, 13, 6, 3, 9, 4, 11, 16, 17], #10c + [8, 14, 7, 4, 10, 5, 9, 15, 17], + [9, 16, 15, 8, 6, 1, 13, 0, 14], + [10, 17, 16, 6, 7, 2, 14, 1, 12], + [11, 15, 17, 7, 8, 2, 13, 0, 12], + [12, 0, 14,11, 17, 9, 16, 3, 4], #15c + [13, 1, 12, 9, 15, 10, 17, 4, 5], + [14, 2, 13,10, 16, 5, 3, 11, 15], + ] + ) class SimpleMesh: _DEFAULT_K_LEVEL = 10 @@ -358,6 +381,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.c2e = SimpleMeshData.c2e_table self.c2e2cO = SimpleMeshData.c2e2cO_table self.c2e2c = SimpleMeshData.c2e2c_table + self.c2e2c2e2c = SimpleMeshData.c2e2c2e2c_table self.e2c2eO = SimpleMeshData.e2c2eO_table self.e2c2e = SimpleMeshData.e2c2e_table self.e2c2v = SimpleMeshData.e2c2v_table @@ -368,6 +392,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_c2e = self.c2e.shape[1] self.n_c2e2cO = self.c2e2cO.shape[1] self.n_c2e2c = self.c2e2c.shape[1] + self.n_c2e2c2e2c = self.c2e2c2e2c.shape[1] self.n_e2c2eO = self.e2c2eO.shape[1] self.n_e2c2e = self.e2c2e.shape[1] self.n_e2c2v = self.e2c2v.shape[1] @@ -385,6 +410,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): C2EDim: self.n_c2e, C2E2CODim: self.n_c2e2cO, C2E2CDim: self.n_c2e2c, + C2E2C2E2CDim: self.n_c2e2c2e2c, E2C2EODim: self.n_e2c2eO, E2C2EDim: self.n_e2c2e, V2CDim: self.n_v2c, @@ -405,6 +431,9 @@ def get_c2e2cO_offset_provider(self) -> NeighborTableOffsetProvider: def get_c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e2c, CellDim, CellDim, self.n_c2e2c) + def get_c2e2c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e2c2e2c, CellDim, CellDim, self.n_c2e2c2e2c) + def get_e2c2eO_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2eO, EdgeDim, EdgeDim, self.n_e2c2eO) From 1d48a14be887fb2aedfb6e5860a6f289f763c8ec Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 27 Jul 2023 11:36:45 +0200 Subject: [PATCH 082/105] Merge origin/main into dsl_advection_52_part2 --- .coveragerc | 45 -- .github/workflows/icon4py-qa.yml | 30 +- .github/workflows/icon4py-tox.yml | 21 +- .github/workflows/icon4pytools-qa.yml | 16 + .github/workflows/icon4pytools-tox.yml | 26 +- .isort.cfg | 15 - .mypy.ini | 29 -- README.md | 40 +- atm_dyn_iconam/README.md | 9 - atm_dyn_iconam/requirements-dev.txt | 3 - atm_dyn_iconam/requirements.txt | 3 - atm_dyn_iconam/setup.cfg | 50 --- ...st_apply_nabla2_and_nabla4_global_to_vn.py | 69 --- .../test_apply_nabla2_and_nabla4_to_vn.py | 79 ---- ..._apply_nabla2_to_vn_in_lateral_boundary.py | 47 -- .../tests/test_apply_nabla2_to_w.py | 66 --- ...pply_nabla2_to_w_in_upper_damping_layer.py | 57 --- ...st_calculate_diagnostics_for_turbulence.py | 63 --- ...ate_horizontal_gradients_for_turbulence.py | 61 --- .../tests/test_calculate_nabla2_for_w.py | 47 -- .../tests/test_calculate_nabla2_for_z.py | 62 --- ...n_coefficient_for_grid_point_cold_pools.py | 47 -- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 50 --- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 59 --- .../test_mo_math_divrot_rot_vertex_ri_dsl.py | 55 --- ...ath_gradients_grad_green_gauss_cell_dsl.py | 83 ---- ...est_mo_solve_nonhydro_4th_order_divdamp.py | 53 --- .../test_mo_solve_nonhydro_stencil_01.py | 49 -- .../test_mo_solve_nonhydro_stencil_02.py | 52 --- .../test_mo_solve_nonhydro_stencil_03.py | 44 -- .../test_mo_solve_nonhydro_stencil_04.py | 55 --- .../test_mo_solve_nonhydro_stencil_06.py | 51 --- .../test_mo_solve_nonhydro_stencil_07.py | 60 --- ...test_mo_solve_nonhydro_stencil_11_lower.py | 43 -- .../test_mo_solve_nonhydro_stencil_12.py | 65 --- .../test_mo_solve_nonhydro_stencil_13.py | 60 --- .../test_mo_solve_nonhydro_stencil_14.py | 47 -- .../test_mo_solve_nonhydro_stencil_15.py | 49 -- .../test_mo_solve_nonhydro_stencil_17.py | 74 --- .../test_mo_solve_nonhydro_stencil_18.py | 60 --- .../test_mo_solve_nonhydro_stencil_19.py | 75 ---- .../test_mo_solve_nonhydro_stencil_22.py | 59 --- .../test_mo_solve_nonhydro_stencil_23.py | 87 ---- .../test_mo_solve_nonhydro_stencil_24.py | 71 --- .../test_mo_solve_nonhydro_stencil_25.py | 49 -- .../test_mo_solve_nonhydro_stencil_26.py | 48 -- .../test_mo_solve_nonhydro_stencil_27.py | 61 --- .../test_mo_solve_nonhydro_stencil_28.py | 48 -- .../test_mo_solve_nonhydro_stencil_29.py | 50 --- .../test_mo_solve_nonhydro_stencil_30.py | 76 ---- .../test_mo_solve_nonhydro_stencil_31.py | 49 -- .../test_mo_solve_nonhydro_stencil_32.py | 64 --- .../test_mo_solve_nonhydro_stencil_33.py | 49 -- .../test_mo_solve_nonhydro_stencil_34.py | 63 --- .../test_mo_solve_nonhydro_stencil_35.py | 52 --- .../test_mo_solve_nonhydro_stencil_37.py | 54 --- .../test_mo_solve_nonhydro_stencil_41.py | 59 --- .../test_mo_solve_nonhydro_stencil_42.py | 96 ---- .../test_mo_solve_nonhydro_stencil_43.py | 80 ---- .../test_mo_solve_nonhydro_stencil_44.py | 89 ---- .../test_mo_solve_nonhydro_stencil_45.py | 40 -- .../test_mo_solve_nonhydro_stencil_46.py | 49 -- .../test_mo_solve_nonhydro_stencil_47.py | 51 --- .../test_mo_solve_nonhydro_stencil_48.py | 101 ----- .../test_mo_solve_nonhydro_stencil_49.py | 101 ----- .../test_mo_solve_nonhydro_stencil_50.py | 64 --- .../test_mo_solve_nonhydro_stencil_53.py | 51 --- .../test_mo_solve_nonhydro_stencil_54.py | 50 --- .../test_mo_solve_nonhydro_stencil_55.py | 124 ----- .../test_mo_solve_nonhydro_stencil_56_63.py | 53 --- .../test_mo_solve_nonhydro_stencil_57.py | 40 -- .../test_mo_solve_nonhydro_stencil_58.py | 68 --- .../test_mo_solve_nonhydro_stencil_59.py | 45 -- .../test_mo_solve_nonhydro_stencil_60.py | 62 --- .../test_mo_solve_nonhydro_stencil_61.py | 78 ---- .../test_mo_solve_nonhydro_stencil_62.py | 51 --- .../test_mo_solve_nonhydro_stencil_64.py | 41 -- .../test_mo_solve_nonhydro_stencil_65.py | 79 ---- .../test_mo_solve_nonhydro_stencil_66.py | 73 --- .../test_mo_solve_nonhydro_stencil_67.py | 58 --- .../test_mo_solve_nonhydro_stencil_68.py | 82 ---- .../test_mo_velocity_advection_stencil_01.py | 51 --- .../test_mo_velocity_advection_stencil_04.py | 50 --- .../test_mo_velocity_advection_stencil_05.py | 58 --- .../test_mo_velocity_advection_stencil_07.py | 91 ---- .../test_mo_velocity_advection_stencil_08.py | 51 --- .../test_mo_velocity_advection_stencil_09.py | 51 --- .../test_mo_velocity_advection_stencil_11.py | 39 -- .../test_mo_velocity_advection_stencil_12.py | 40 -- .../test_mo_velocity_advection_stencil_13.py | 47 -- .../test_mo_velocity_advection_stencil_14.py | 90 ---- .../test_mo_velocity_advection_stencil_15.py | 49 -- .../test_mo_velocity_advection_stencil_17.py | 49 -- .../test_mo_velocity_advection_stencil_18.py | 112 ----- ...d_for_grid_point_cold_pools_enhancement.py | 58 --- ...orary_fields_for_turbulence_diagnostics.py | 80 ---- .../tests/test_update_theta_and_exner.py | 62 --- atm_dyn_iconam/tests/test_utils/__init__.py | 12 - base-requirements-dev.txt | 2 +- common/README.md | 9 - common/requirements-dev.txt | 2 - common/requirements.txt | 2 - common/setup.cfg | 50 --- common/src/icon4py/common/__init__.py | 12 - jenkins/spack-PR | 10 - jenkins/spack-PR-icon | 95 ++++ jenkins/spack-PR-stable | 42 ++ .flake8 => model/.flake8 | 3 - .../.license_header.txt | 0 .../.pre-commit-config.yaml | 14 +- LICENSE => model/LICENSE | 0 model/README.md | 24 + model/atmosphere/dycore/.bumpversion.cfg | 10 + model/atmosphere/dycore/README.md | 9 + model/atmosphere/dycore/pyproject.toml | 118 +++++ model/atmosphere/dycore/requirements-dev.txt | 3 + model/atmosphere/dycore/requirements.txt | 3 + .../model/atmosphere/dycore/__init__.py | 21 +- .../apply_diffusion_to_theta_and_exner.py | 14 +- .../dycore}/apply_diffusion_to_vn.py | 15 +- ...ute_horizontal_gradients_for_turbulance.py | 21 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 9 +- .../dycore}/apply_nabla2_and_nabla4_to_vn.py | 5 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 9 +- .../atmosphere/dycore}/apply_nabla2_to_w.py | 5 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 9 +- ...te_diagnostic_quantities_for_turbulence.py | 6 +- .../calculate_diagnostics_for_turbulence.py | 10 +- ..._coefficients_for_grid_point_cold_pools.py | 10 +- ...ate_horizontal_gradients_for_turbulence.py | 9 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 23 +- .../dycore}/calculate_nabla2_for_theta.py | 12 +- .../dycore}/calculate_nabla2_for_w.py | 5 +- .../dycore}/calculate_nabla2_for_z.py | 5 +- .../dycore}/calculate_nabla2_of_theta.py | 5 +- .../atmosphere/dycore}/calculate_nabla4.py | 10 +- ...n_coefficient_for_grid_point_cold_pools.py | 9 +- .../mo_advection_traj_btraj_compute_o1_dsl.py | 11 +- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 9 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 9 +- .../mo_math_divrot_rot_vertex_ri_dsl.py | 5 +- ...ath_gradients_grad_green_gauss_cell_dsl.py | 5 +- .../mo_solve_nonhydro_4th_order_divdamp.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_01.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_02.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_03.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_04.py | 7 +- .../dycore}/mo_solve_nonhydro_stencil_05.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_06.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_07.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_08.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_09.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_10.py | 9 +- .../mo_solve_nonhydro_stencil_11_lower.py | 5 +- .../mo_solve_nonhydro_stencil_11_upper.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_12.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_13.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_14.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_15.py | 5 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 18 +- .../dycore}/mo_solve_nonhydro_stencil_17.py | 10 +- .../dycore}/mo_solve_nonhydro_stencil_18.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_19.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_20.py | 16 +- .../dycore}/mo_solve_nonhydro_stencil_21.py | 10 +- .../dycore}/mo_solve_nonhydro_stencil_22.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_23.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_24.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_25.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_26.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_27.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_28.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_29.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_30.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_31.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_32.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_33.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_34.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_35.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_36.py | 8 +- .../dycore}/mo_solve_nonhydro_stencil_37.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_38.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_39.py | 16 +- .../dycore}/mo_solve_nonhydro_stencil_40.py | 24 +- .../dycore}/mo_solve_nonhydro_stencil_41.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_42.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_43.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_44.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_45.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_45_b.py | 7 +- .../dycore}/mo_solve_nonhydro_stencil_46.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_47.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_48.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_49.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_50.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_51.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_52.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_53.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_54.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_55.py | 18 +- .../mo_solve_nonhydro_stencil_56_63.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_57.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_58.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_59.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_60.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_61.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_62.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_64.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_65.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_66.py | 5 +- .../dycore}/mo_solve_nonhydro_stencil_67.py | 9 +- .../dycore}/mo_solve_nonhydro_stencil_68.py | 10 +- .../mo_velocity_advection_stencil_01.py | 5 +- .../mo_velocity_advection_stencil_02.py | 5 +- .../mo_velocity_advection_stencil_03.py | 5 +- .../mo_velocity_advection_stencil_04.py | 9 +- .../mo_velocity_advection_stencil_05.py | 5 +- .../mo_velocity_advection_stencil_06.py | 5 +- .../mo_velocity_advection_stencil_07.py | 16 +- .../mo_velocity_advection_stencil_08.py | 5 +- .../mo_velocity_advection_stencil_09.py | 5 +- .../mo_velocity_advection_stencil_10.py | 9 +- .../mo_velocity_advection_stencil_11.py | 5 +- .../mo_velocity_advection_stencil_12.py | 5 +- .../mo_velocity_advection_stencil_13.py | 5 +- .../mo_velocity_advection_stencil_14.py | 13 +- .../mo_velocity_advection_stencil_15.py | 5 +- .../mo_velocity_advection_stencil_16.py | 13 +- .../mo_velocity_advection_stencil_17.py | 5 +- .../mo_velocity_advection_stencil_18.py | 8 +- .../mo_velocity_advection_stencil_19.py | 10 +- .../mo_velocity_advection_stencil_20.py | 9 +- .../icon4py/model/atmosphere/dycore}/py.typed | 0 ...d_for_grid_point_cold_pools_enhancement.py | 5 +- ...orary_fields_for_turbulence_diagnostics.py | 5 +- ...fusion_nabla_of_theta_over_steep_points.py | 25 +- .../dycore}/update_theta_and_exner.py | 9 +- .../atmosphere/dycore/tests}/__init__.py | 0 .../atmosphere/dycore/tests/conftest.py | 21 + ...st_apply_nabla2_and_nabla4_global_to_vn.py | 51 +++ .../test_apply_nabla2_and_nabla4_to_vn.py | 69 +++ ..._apply_nabla2_to_vn_in_lateral_boundary.py | 47 ++ .../dycore/tests/test_apply_nabla2_to_w.py | 54 +++ ...pply_nabla2_to_w_in_upper_damping_layer.py | 52 +++ ...st_calculate_diagnostics_for_turbulence.py | 49 ++ ...ate_horizontal_gradients_for_turbulence.py | 57 +++ ...ate_nabla2_and_smag_coefficients_for_vn.py | 19 +- .../tests/test_calculate_nabla2_for_w.py | 48 ++ .../tests/test_calculate_nabla2_for_z.py | 56 +++ .../tests/test_calculate_nabla2_of_theta.py | 17 +- .../dycore}/tests/test_calculate_nabla4.py | 19 +- ...n_coefficient_for_grid_point_cold_pools.py | 45 ++ ..._mo_advection_traj_btraj_compute_o1_dsl.py | 31 +- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 48 ++ ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 58 +++ .../test_mo_math_divrot_rot_vertex_ri_dsl.py | 48 ++ ...ath_gradients_grad_green_gauss_cell_dsl.py | 77 ++++ ...est_mo_solve_nonhydro_4th_order_divdamp.py | 49 ++ .../test_mo_solve_nonhydro_stencil_01.py | 43 ++ .../test_mo_solve_nonhydro_stencil_02.py | 59 +++ .../test_mo_solve_nonhydro_stencil_03.py | 39 ++ .../test_mo_solve_nonhydro_stencil_04.py | 56 +++ .../test_mo_solve_nonhydro_stencil_05.py | 9 +- .../test_mo_solve_nonhydro_stencil_06.py | 47 ++ .../test_mo_solve_nonhydro_stencil_07.py | 61 +++ .../test_mo_solve_nonhydro_stencil_08.py | 9 +- .../test_mo_solve_nonhydro_stencil_09.py | 9 +- .../test_mo_solve_nonhydro_stencil_10.py | 9 +- ...test_mo_solve_nonhydro_stencil_11_lower.py | 38 ++ ...test_mo_solve_nonhydro_stencil_11_upper.py | 9 +- .../test_mo_solve_nonhydro_stencil_12.py | 63 +++ .../test_mo_solve_nonhydro_stencil_13.py | 61 +++ .../test_mo_solve_nonhydro_stencil_14.py | 42 ++ .../test_mo_solve_nonhydro_stencil_15.py | 42 ++ ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 36 +- .../test_mo_solve_nonhydro_stencil_17.py | 64 +++ .../test_mo_solve_nonhydro_stencil_18.py | 50 +++ .../test_mo_solve_nonhydro_stencil_19.py | 65 +++ .../test_mo_solve_nonhydro_stencil_20.py | 30 +- .../test_mo_solve_nonhydro_stencil_21.py | 26 +- .../test_mo_solve_nonhydro_stencil_22.py | 59 +++ .../test_mo_solve_nonhydro_stencil_23.py | 81 ++++ .../test_mo_solve_nonhydro_stencil_24.py | 68 +++ .../test_mo_solve_nonhydro_stencil_25.py | 48 ++ .../test_mo_solve_nonhydro_stencil_26.py | 43 ++ .../test_mo_solve_nonhydro_stencil_27.py | 56 +++ .../test_mo_solve_nonhydro_stencil_28.py | 43 ++ .../test_mo_solve_nonhydro_stencil_29.py | 49 ++ .../test_mo_solve_nonhydro_stencil_30.py | 67 +++ .../test_mo_solve_nonhydro_stencil_31.py | 48 ++ .../test_mo_solve_nonhydro_stencil_32.py | 61 +++ .../test_mo_solve_nonhydro_stencil_33.py | 42 ++ .../test_mo_solve_nonhydro_stencil_34.py | 56 +++ .../test_mo_solve_nonhydro_stencil_35.py | 58 +++ .../test_mo_solve_nonhydro_stencil_36.py | 9 +- .../test_mo_solve_nonhydro_stencil_37.py | 53 +++ .../test_mo_solve_nonhydro_stencil_38.py | 9 +- .../test_mo_solve_nonhydro_stencil_39.py | 9 +- .../test_mo_solve_nonhydro_stencil_40.py | 9 +- .../test_mo_solve_nonhydro_stencil_41.py | 59 +++ .../test_mo_solve_nonhydro_stencil_42.py | 85 ++++ .../test_mo_solve_nonhydro_stencil_43.py | 74 +++ .../test_mo_solve_nonhydro_stencil_44.py | 81 ++++ .../test_mo_solve_nonhydro_stencil_45.py | 39 ++ .../test_mo_solve_nonhydro_stencil_46.py | 42 ++ .../test_mo_solve_nonhydro_stencil_47.py | 48 ++ .../test_mo_solve_nonhydro_stencil_48.py | 92 ++++ .../test_mo_solve_nonhydro_stencil_49.py | 92 ++++ .../test_mo_solve_nonhydro_stencil_50.py | 56 +++ .../test_mo_solve_nonhydro_stencil_51.py | 14 +- .../test_mo_solve_nonhydro_stencil_52.py | 17 +- .../test_mo_solve_nonhydro_stencil_53.py | 43 ++ .../test_mo_solve_nonhydro_stencil_54.py | 45 ++ .../test_mo_solve_nonhydro_stencil_55.py | 107 +++++ .../test_mo_solve_nonhydro_stencil_56_63.py | 49 ++ .../test_mo_solve_nonhydro_stencil_57.py | 39 ++ .../test_mo_solve_nonhydro_stencil_58.py | 59 +++ .../test_mo_solve_nonhydro_stencil_59.py | 45 ++ .../test_mo_solve_nonhydro_stencil_60.py | 54 +++ .../test_mo_solve_nonhydro_stencil_61.py | 73 +++ .../test_mo_solve_nonhydro_stencil_62.py | 49 ++ .../test_mo_solve_nonhydro_stencil_64.py | 39 ++ .../test_mo_solve_nonhydro_stencil_65.py | 68 +++ .../test_mo_solve_nonhydro_stencil_66.py | 68 +++ .../test_mo_solve_nonhydro_stencil_67.py | 55 +++ .../test_mo_solve_nonhydro_stencil_68.py | 74 +++ .../test_mo_velocity_advection_stencil_01.py | 48 ++ .../test_mo_velocity_advection_stencil_02.py | 17 +- .../test_mo_velocity_advection_stencil_03.py | 17 +- .../test_mo_velocity_advection_stencil_04.py | 58 +++ .../test_mo_velocity_advection_stencil_05.py | 54 +++ .../test_mo_velocity_advection_stencil_06.py | 17 +- .../test_mo_velocity_advection_stencil_07.py | 80 ++++ .../test_mo_velocity_advection_stencil_08.py | 48 ++ .../test_mo_velocity_advection_stencil_09.py | 48 ++ .../test_mo_velocity_advection_stencil_10.py | 13 +- .../test_mo_velocity_advection_stencil_11.py | 41 ++ .../test_mo_velocity_advection_stencil_12.py | 39 ++ .../test_mo_velocity_advection_stencil_13.py | 41 ++ .../test_mo_velocity_advection_stencil_14.py | 86 ++++ .../test_mo_velocity_advection_stencil_15.py | 46 ++ .../test_mo_velocity_advection_stencil_16.py | 9 +- .../test_mo_velocity_advection_stencil_17.py | 46 ++ .../test_mo_velocity_advection_stencil_18.py | 100 +++++ .../test_mo_velocity_advection_stencil_19.py | 18 +- .../test_mo_velocity_advection_stencil_20.py | 9 +- ...d_for_grid_point_cold_pools_enhancement.py | 59 +++ ...orary_fields_for_turbulence_diagnostics.py | 73 +++ ...fusion_nabla_of_theta_over_steep_points.py | 25 +- .../tests/test_update_theta_and_exner.py | 58 +++ model/common/.bumpversion.cfg | 10 + model/common/README.md | 9 + model/common/pyproject.toml | 117 +++++ model/common/requirements-dev.txt | 2 + model/common/requirements.txt | 2 + .../src/icon4py/model/common/__init__.py | 21 +- .../src/icon4py/model}/common/dimension.py | 0 .../common/src/icon4py/model/common}/py.typed | 0 .../model/common/test_utils}/__init__.py | 0 .../model/common}/test_utils/helpers.py | 75 +++- .../model/common}/test_utils/simple_mesh.py | 17 +- model/requirements-dev.txt | 3 + model/requirements.txt | 3 + model/tox.ini | 33 ++ pytest.ini | 1 + requirements-dev.txt | 8 +- requirements.txt | 8 +- spack/gt4py-stable/spack.yaml | 5 +- tools/.bumpversion.cfg | 2 +- tools/.license_header.txt | 12 + tools/.pre-commit-config.yaml | 5 +- tools/pyproject.toml | 4 +- tools/requirements-dev.txt | 6 +- tools/requirements.txt | 6 +- tools/src/icon4pytools/__init__.py | 2 +- tools/src/icon4pytools/common/__init__.py | 4 + tools/src/icon4pytools/f2ser/parse.py | 156 +++++-- tools/src/icon4pytools/icon4pygen/backend.py | 2 +- .../icon4pygen/bindings/entities.py | 1 - .../icon4pygen/bindings/locations.py | 3 +- tools/src/icon4pytools/icon4pygen/cli.py | 7 +- .../icon4pytools/icon4pygen/icochainsize.py | 3 +- tools/src/icon4pytools/icon4pygen/metadata.py | 3 +- .../liskov/codegen/integration/deserialise.py | 2 - .../src/icon4pytools/liskov/external/gt4py.py | 13 +- tools/src/icon4pytools/liskov/py.typed | 0 .../fortran_samples/multiline_example.f90 | 423 ++++++++++++++++++ tools/tests/f2ser/test_parsing.py | 13 + tools/tests/icon4pygen/helpers.py | 4 +- tools/tests/icon4pygen/test_codegen.py | 12 +- tools/tests/icon4pygen/test_exceptions.py | 2 +- .../tests/icon4pygen/test_field_rendering.py | 3 +- tools/tests/icon4pygen/test_metadata.py | 2 +- tools/tests/liskov/test_external.py | 2 +- tools/tox.ini | 2 + tox.ini | 13 +- 396 files changed, 6834 insertions(+), 6310 deletions(-) delete mode 100644 .coveragerc delete mode 100644 .isort.cfg delete mode 100644 .mypy.ini delete mode 100644 atm_dyn_iconam/README.md delete mode 100644 atm_dyn_iconam/requirements-dev.txt delete mode 100644 atm_dyn_iconam/requirements.txt delete mode 100644 atm_dyn_iconam/setup.cfg delete mode 100644 atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_global_to_vn.py delete mode 100644 atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_to_vn.py delete mode 100644 atm_dyn_iconam/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py delete mode 100644 atm_dyn_iconam/tests/test_apply_nabla2_to_w.py delete mode 100644 atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py delete mode 100644 atm_dyn_iconam/tests/test_calculate_diagnostics_for_turbulence.py delete mode 100644 atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py delete mode 100644 atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py delete mode 100644 atm_dyn_iconam/tests/test_calculate_nabla2_for_z.py delete mode 100644 atm_dyn_iconam/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py delete mode 100644 atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py delete mode 100644 atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py delete mode 100644 atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py delete mode 100644 atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_4th_order_divdamp.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_01.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_02.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_03.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_04.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_06.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_07.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_lower.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_12.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_13.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_14.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_15.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_17.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_18.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_22.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_23.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_24.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_26.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_27.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_28.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_29.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_32.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_33.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_34.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_35.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_37.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_42.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_43.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_44.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_45.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_46.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_47.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_48.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_49.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_50.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_53.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_54.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_55.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_56_63.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_57.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_58.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_59.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_60.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_61.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_62.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_64.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_65.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_66.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_67.py delete mode 100644 atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_68.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_04.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_05.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_07.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_11.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_12.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_13.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_14.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_15.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py delete mode 100644 atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py delete mode 100644 atm_dyn_iconam/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py delete mode 100644 atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py delete mode 100644 atm_dyn_iconam/tests/test_update_theta_and_exner.py delete mode 100644 atm_dyn_iconam/tests/test_utils/__init__.py delete mode 100644 common/README.md delete mode 100644 common/requirements-dev.txt delete mode 100644 common/requirements.txt delete mode 100644 common/setup.cfg delete mode 100644 common/src/icon4py/common/__init__.py create mode 100644 jenkins/spack-PR-icon create mode 100644 jenkins/spack-PR-stable rename .flake8 => model/.flake8 (91%) rename .license_header.txt => model/.license_header.txt (100%) rename .pre-commit-config.yaml => model/.pre-commit-config.yaml (90%) rename LICENSE => model/LICENSE (100%) create mode 100644 model/README.md create mode 100644 model/atmosphere/dycore/.bumpversion.cfg create mode 100644 model/atmosphere/dycore/README.md create mode 100644 model/atmosphere/dycore/pyproject.toml create mode 100644 model/atmosphere/dycore/requirements-dev.txt create mode 100644 model/atmosphere/dycore/requirements.txt rename common/setup.py => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py (55%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_diffusion_to_theta_and_exner.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_diffusion_to_vn.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py (83%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_nabla2_and_nabla4_global_to_vn.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_nabla2_and_nabla4_to_vn.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_nabla2_to_vn_in_lateral_boundary.py (83%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_nabla2_to_w.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/apply_nabla2_to_w_in_upper_damping_layer.py (83%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_diagnostic_quantities_for_turbulence.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_diagnostics_for_turbulence.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_horizontal_gradients_for_turbulence.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla2_and_smag_coefficients_for_vn.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla2_for_theta.py (80%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla2_for_w.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla2_for_z.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla2_of_theta.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/calculate_nabla4.py (92%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/enhance_diffusion_coefficient_for_grid_point_cold_pools.py (81%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_advection_traj_btraj_compute_o1_dsl.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py (82%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_intp_rbf_rbf_vec_interpol_vertex.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_math_divrot_rot_vertex_ri_dsl.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_math_gradients_grad_green_gauss_cell_dsl.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_4th_order_divdamp.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_01.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_02.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_03.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_04.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_05.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_06.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_07.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_08.py (92%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_09.py (94%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_10.py (95%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_11_lower.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_11_upper.py (92%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_12.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_13.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_14.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_15.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py (94%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_17.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_18.py (82%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_19.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_20.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_21.py (92%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_22.py (85%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_23.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_24.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_25.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_26.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_27.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_28.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_29.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_30.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_31.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_32.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_33.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_34.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_35.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_36.py (85%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_37.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_38.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_39.py (83%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_40.py (76%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_41.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_42.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_43.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_44.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_45.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_45_b.py (80%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_46.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_47.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_48.py (94%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_49.py (94%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_50.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_51.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_52.py (95%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_53.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_54.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_55.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_56_63.py (83%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_57.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_58.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_59.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_60.py (85%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_61.py (93%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_62.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_64.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_65.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_66.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_67.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_solve_nonhydro_stencil_68.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_01.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_02.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_03.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_04.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_05.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_06.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_07.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_08.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_09.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_10.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_11.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_12.py (87%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_13.py (89%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_14.py (86%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_15.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_16.py (78%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_17.py (88%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_18.py (90%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_19.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/mo_velocity_advection_stencil_20.py (93%) rename {common/src/icon4py/common => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/py.typed (100%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/temporary_field_for_grid_point_cold_pools_enhancement.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/temporary_fields_for_turbulence_diagnostics.py (91%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (84%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/update_theta_and_exner.py (85%) rename {atm_dyn_iconam => model/atmosphere/dycore/tests}/__init__.py (100%) rename atm_dyn_iconam/tests/__init__.py => model/atmosphere/dycore/tests/conftest.py (54%) create mode 100644 model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py create mode 100644 model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py create mode 100644 model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py create mode 100644 model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py create mode 100644 model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py create mode 100644 model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py create mode 100644 model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py (94%) create mode 100644 model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py create mode 100644 model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_calculate_nabla2_of_theta.py (77%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_calculate_nabla4.py (90%) create mode 100644 model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py (83%) create mode 100644 model/atmosphere/dycore/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py create mode 100644 model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py create mode 100644 model/atmosphere/dycore/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py create mode 100644 model/atmosphere/dycore/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_4th_order_divdamp.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_01.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_03.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_05.py (83%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_07.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_08.py (88%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_09.py (91%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_10.py (95%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_lower.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_11_upper.py (87%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_12.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_13.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_14.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py (87%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_17.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_18.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_20.py (85%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_21.py (89%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_22.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_23.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_24.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_27.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_30.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_31.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_32.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_33.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_34.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_35.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_36.py (87%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_37.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_38.py (84%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_39.py (86%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_40.py (88%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_41.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_43.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_45.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_46.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_48.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_49.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_50.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_51.py (89%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_solve_nonhydro_stencil_52.py (87%) create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_53.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_56_63.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_57.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_59.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_61.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_64.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_66.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_67.py create mode 100644 model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_01.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_02.py (79%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_03.py (70%) create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_04.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_05.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_06.py (76%) create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_07.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_10.py (78%) create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_11.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_12.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_15.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_16.py (86%) create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_17.py create mode 100644 model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_19.py (90%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_mo_velocity_advection_stencil_20.py (94%) create mode 100644 model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py create mode 100644 model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (86%) create mode 100644 model/atmosphere/dycore/tests/test_update_theta_and_exner.py create mode 100644 model/common/.bumpversion.cfg create mode 100644 model/common/README.md create mode 100644 model/common/pyproject.toml create mode 100644 model/common/requirements-dev.txt create mode 100644 model/common/requirements.txt rename atm_dyn_iconam/setup.py => model/common/src/icon4py/model/common/__init__.py (55%) rename {common/src/icon4py => model/common/src/icon4py/model}/common/dimension.py (100%) rename {tools/src/icon4pytools/icon4pygen => model/common/src/icon4py/model/common}/py.typed (100%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/common/src/icon4py/model/common/test_utils}/__init__.py (100%) rename {atm_dyn_iconam/tests => model/common/src/icon4py/model/common}/test_utils/helpers.py (58%) rename {atm_dyn_iconam/tests => model/common/src/icon4py/model/common}/test_utils/simple_mesh.py (94%) create mode 100644 model/requirements-dev.txt create mode 100644 model/requirements.txt create mode 100644 model/tox.ini create mode 100644 tools/.license_header.txt delete mode 100644 tools/src/icon4pytools/liskov/py.typed create mode 100644 tools/tests/f2ser/fortran_samples/multiline_example.f90 diff --git a/.coveragerc b/.coveragerc deleted file mode 100644 index 2c097ba2c0..0000000000 --- a/.coveragerc +++ /dev/null @@ -1,45 +0,0 @@ -# Coverage configuration -# Reference documentation: https://coverage.readthedocs.io/ - -[run] -branch = True -source_pkgs = - icon4py - -[paths] -atm_dyn_iconam_sources = - atm_dyn_iconam/src/icon4py/atm_dyn_iconam/ - .tox/py*/lib/python3.*/site-packages/icon4py/atm_dyn_iconam/ - -advection_sources = - advection/src/icon4py/advection/ - .tox/py*/lib/python3.*/site-packages/icon4py/advection/ - -common_sources = - common/src/icon4py/common/ - .tox/py*/lib/python3.*/site-packages/icon4py/common/ - -pyutils_sources = - pyutils/src/icon4py/pyutils/ - .tox/py*/lib/python3.*/site-packages/icon4py/pyutils/ - -testutils_sources = - testutils/src/icon4py/testutils/ - .tox/py*/lib/python3.*/site-packages/icon4py/testutils/ - -[report] -ignore_errors = True -show_missing = True - -# Regexes for lines to exclude from consideration -exclude_lines = - # Don't complain if tests don't hit defensive assertion code: - raise AssertionError - raise NotImplementedError - - # Don't complain if non-runnable code isn't run: - if 0: - if __name__ == .__main__.: - -[html] -directory = _reports/coverage_html diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 81584064ff..65cec9d877 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -1,6 +1,7 @@ name: ICON4Py Quality Assurance on: + workflow_dispatch: push: branches: - main @@ -9,6 +10,15 @@ on: pull_request: branches: - main + paths-ignore: + - "tools/**" + types: [opened, reopened, assigned, ready_for_review] + pull_request_review: + branches: + - main + paths-ignore: + - "tools/**" + types: [submitted] jobs: pre-commit: @@ -19,16 +29,18 @@ jobs: uses: actions/setup-python@v3 with: python-version: "3.10" - - name: Install dependencies + - name: Install icon4py-atmosphere-dycore + working-directory: model/atmosphere/dycore + run: | + python -m pip install --upgrade pip setuptools wheel + python -m pip install -r ./requirements-dev.txt + + - name: Install icon4py-common + working-directory: model/common run: | python -m pip install --upgrade pip setuptools wheel - python -m pip install $(cat ./base-requirements-dev.txt | grep pre-commit) - pre-commit install-hooks - # mypy hook uses local repo and thus requires installing icon4py & gt4py - python -m pip install $(cat ./base-requirements-dev.txt | grep mypy) - python -m pip install $(cat ./base-requirements-dev.txt | grep gt4py) - python -m pip install $(cat ./requirements-dev.txt | grep '\-e') - - name: Run checks + python -m pip install -r ./requirements-dev.txt + - name: Run checks icon4py-model run: | python -m pip list - pre-commit run --config .pre-commit-config.yaml --all-files + pre-commit run --config model/.pre-commit-config.yaml --all-files diff --git a/.github/workflows/icon4py-tox.yml b/.github/workflows/icon4py-tox.yml index f39655c350..4d2cdab865 100644 --- a/.github/workflows/icon4py-tox.yml +++ b/.github/workflows/icon4py-tox.yml @@ -1,6 +1,7 @@ name: ICON4Py Tox on: + workflow_dispatch: push: branches: - main @@ -9,7 +10,15 @@ on: pull_request: branches: - main - + paths-ignore: + - "tools/**" + types: [opened, reopened, assigned, ready_for_review] + pull_request_review: + branches: + - main + paths-ignore: + - "tools/**" + types: [submitted] jobs: tox-tests: runs-on: ubuntu-latest @@ -34,15 +43,9 @@ jobs: run: | python -m pip install --upgrade pip setuptools wheel - - name: Test with tox + - name: Test icon4py-model with tox run: | pyversion_no_dot="${{ matrix.python-version }}" pyversion_no_dot="${pyversion_no_dot/./}" pip install tox clang-format - tox -r -e py${pyversion_no_dot} - - - name: Archive code coverage results - uses: actions/upload-artifact@v2 - with: - name: code-coverage-report - path: _reports/coverage_html/ + tox -r -e py${pyversion_no_dot} -c model/ diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index f82deb0295..9afe586f30 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -1,15 +1,30 @@ name: ICON4PyTools Quality Assurance on: + workflow_dispatch: push: branches: - main paths: - "tools/**" - "base-requirements*" + - ".github/workflows/icon4pytools*" pull_request: branches: - main + paths: + - "tools/**" + - "base-requirements*" + - ".github/workflows/icon4pytools*" + types: [opened, reopened, assigned, ready_for_review] + pull_request_review: + branches: + - main + paths: + - "tools/**" + - "base-requirements*" + - ".github/workflows/icon4pytools*" + types: [submitted] jobs: pre-commit: @@ -21,6 +36,7 @@ jobs: with: python-version: "3.10" - name: Install dependencies + working-directory: tools run: | python -m pip install --upgrade pip setuptools wheel python -m pip install -r ./requirements-dev.txt diff --git a/.github/workflows/icon4pytools-tox.yml b/.github/workflows/icon4pytools-tox.yml index 2b68e80181..5dfc6d6d99 100644 --- a/.github/workflows/icon4pytools-tox.yml +++ b/.github/workflows/icon4pytools-tox.yml @@ -1,15 +1,30 @@ name: ICON4PyTools Tox on: + workflow_dispatch: push: branches: - main paths: - "tools/**" - "base-requirements*" + - ".github/workflows/icon4pytools*" pull_request: branches: - main + paths: + - "tools/**" + - "base-requirements*" + - ".github/workflows/icon4pytools*" + types: [opened, reopened, assigned, ready_for_review] + pull_request_review: + branches: + - main + paths: + - "tools/**" + - "base-requirements*" + - ".github/workflows/icon4pytools*" + types: [submitted] jobs: tox-tests: @@ -17,9 +32,6 @@ jobs: strategy: matrix: python-version: ["3.10"] - defaults: - run: - working-directory: tools steps: - uses: actions/checkout@v2 @@ -43,10 +55,4 @@ jobs: pyversion_no_dot="${{ matrix.python-version }}" pyversion_no_dot="${pyversion_no_dot/./}" pip install tox clang-format - tox -r -e py${pyversion_no_dot} - - - name: Archive code coverage results - uses: actions/upload-artifact@v2 - with: - name: code-coverage-report - path: _reports/coverage_html/ + tox -r -e py${pyversion_no_dot} -c tools/ diff --git a/.isort.cfg b/.isort.cfg deleted file mode 100644 index d6a590ce99..0000000000 --- a/.isort.cfg +++ /dev/null @@ -1,15 +0,0 @@ -[settings] -skip_glob = - *.venv/** - _local/** -line_length = 100 -force_grid_wrap = 0 -include_trailing_comma = true -multi_line_output = 3 -use_parentheses = true -lines_after_imports = 2 -default_section = THIRDPARTY -sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER -known_first_party = icon4py -known_third_party = gt4py -profile = black diff --git a/.mypy.ini b/.mypy.ini deleted file mode 100644 index a8e31eb26d..0000000000 --- a/.mypy.ini +++ /dev/null @@ -1,29 +0,0 @@ -[mypy] -# Global options -install_types = True -non_interactive = True - -# File paths and patterns -exclude = (?x)( - /_external_src/.* | - /build/.* | - /setup\.py$ | - /tests/.* | - ) -namespace_packages = True - -# Check options -disallow_incomplete_defs = True -disallow_untyped_defs = True -ignore_missing_imports = False -implicit_reexport = True - -# Warnings -warn_unused_configs = True -warn_unused_ignores = True -warn_redundant_casts = True - -# Formatting -#pretty = True -show_column_numbers = True -show_error_codes = True diff --git a/README.md b/README.md index bbba9acf5f..2b4738ba6d 100644 --- a/README.md +++ b/README.md @@ -2,13 +2,11 @@ # ICON4Py -## Description +ICON4Py hosts Python implementations of various components from the ICON climate and weather model. Additionally, it includes icon4pytools, a collection of command-line interfaces (CLIs), and utilities required for the integration of ICON4Py code into the ICON Fortran model. ICON4Py leverages [GT4Py](https://github.com/GridTools/gt4py) to ensure efficient and performance portable implementations of these components. -ICON4Py contains Python (GT4Py) implementations of ICON (inspired) components for weather and climate models as well as `icon4pytools`, a package containing CLIs and utilties needed for the integration of ICON4Py code into the ICON Fortran model. +## Project Structure -## Project structure - -Each directory contains Python packages of ICON components or utility packages that are deployable on their own. As these packages are not available from a package repository (yet), location of dependencies within this repository have to be provided explicitly, e.g. by installing the dependencies first. See [Installation instructions](#installation-instructions). +The repository is organized into directories, each containing independent Python namespace packages for different ICON components or utility packages. These packages can be installed independently. Since these packages are not available from a package repository (yet), you need to specify the location of dependencies within this repository. This can be done by installing the required dependencies first. Refer to the [Installation instructions](#installation-instructions) below. ## Installation instructions @@ -49,7 +47,7 @@ pip install --src _external_src -r requirements-dev.txt pytest -v ``` -The `--src _external_src` option tells `pip` to use a specific folder as base path for checked out sources, which is very convenient for development tasks involving changes in external dependencies like `gt4py`. For convenience, `./_external_src` has been already added to the repository `.gitignore`. +The `--src _external_src` option tells `pip` to use a specific folder as the base path for checked out sources, which is very convenient for development tasks involving changes in external dependencies like `gt4py`. For convenience, `./_external_src` has been already added to the repository `.gitignore`. ### Installation of specific subpackages @@ -70,7 +68,7 @@ source .venv/bin/activate pip install --upgrade wheel # Install a specific ICON4Py subpackage and its dependencies -cd _SUBPACKAGE_ # where _SUBPACKAGE_ in atm_dyn_iconam | tools | ... +cd _SUBPACKAGE_ # where _SUBPACKAGE_ in model/atmosphere/dycore | tools | ... pip install -r requirements-dev.txt # or in the case of there being a pyproject.toml file @@ -83,24 +81,14 @@ After following the installation instructions above using the development requir ### Code quality checks -[pre-commit](https://pre-commit.com/) is used to run several linting and checking tools. It should always be executed locally before opening a pull request. `pre-commit` can also be installed as a _git hook_ to automatically check the staged changes before committing: +[pre-commit](https://pre-commit.com/) is used to run several linting and checking tools. It should always be executed locally before opening a pull request. When executing pre-commit locally you can either run it for the `model` or `tools` folder: -```bash -# Install pre-commit as a git hook and set up all the tools -pre-commit install --install-hooks -``` - -Or it can be executed on demand from the command line: +For example to run code checks on all components in `icon4py.model` you can do: ```bash -# Check only the staged changes -pre-commit run - -# Check all the files in the repository -pre-commit run -a - -# Run only some of the tools (e.g. mypy) -pre-commit run -a mypy +# running precommit for all components in model +cd model/ +pre-commit run --all-files ``` ### Testing @@ -126,3 +114,11 @@ tox -e py310 ``` The default `tox` environment is configured to generate HTML test coverage reports in `_reports/coverage_html/`. + +### Benchmarking + +We use [`pytest-benchmark`](https://pytest-benchmark.readthedocs.io/en/latest/) to benchmark the execution time of stencils in icon4py. To disable benchmarking during testing you can use `--benchmark-disable` when invoking `pytest`. + +### More Information + +For more information please consult the package specific READMEs found in the `model` and `tools` folders. diff --git a/atm_dyn_iconam/README.md b/atm_dyn_iconam/README.md deleted file mode 100644 index d2936ecd28..0000000000 --- a/atm_dyn_iconam/README.md +++ /dev/null @@ -1,9 +0,0 @@ -# icon4py-atm_dyn_iconam - -## Description - -Contains code ported from ICON `src/atm_dyn_iconam`. - -## Installation instructions - -Check `README.md` file in the root of the repository. diff --git a/atm_dyn_iconam/requirements-dev.txt b/atm_dyn_iconam/requirements-dev.txt deleted file mode 100644 index 2cb7bda5e8..0000000000 --- a/atm_dyn_iconam/requirements-dev.txt +++ /dev/null @@ -1,3 +0,0 @@ --r ../base-requirements-dev.txt --e ../common --e . diff --git a/atm_dyn_iconam/requirements.txt b/atm_dyn_iconam/requirements.txt deleted file mode 100644 index c8420e5bf8..0000000000 --- a/atm_dyn_iconam/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ --r ../base-requirements.txt -../common -. diff --git a/atm_dyn_iconam/setup.cfg b/atm_dyn_iconam/setup.cfg deleted file mode 100644 index a8965f6db5..0000000000 --- a/atm_dyn_iconam/setup.cfg +++ /dev/null @@ -1,50 +0,0 @@ -# This file is mainly used to configure package creation with setuptools. -# Documentation: -# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files -# -[metadata] -name = icon4py_atm_dyn_iconam -description = Icon inspired code in Python and GT4Py -long_description = file: README.md -long_description_content_type = text/markdown -url = https://github.com/C2SM/icon4py -author = ETH Zurich -author_email = gridtools@cscs.ch -license = gpl3 -license_files = LICENSE -platforms = Linux, Mac -classifiers = - Development Status :: 3 - Alpha - Intended Audience :: Science/Research - License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) - Operating System :: POSIX - Programming Language :: Python - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.10 - Programming Language :: Python :: Implementation :: CPython - Topic :: Scientific/Engineering :: Atmospheric Science - Topic :: Scientific/Engineering :: Mathematics - Topic :: Scientific/Engineering :: Physics -project_urls = - Source Code = https://github.com/C2SM/icon4py - -[options] -packages = find_namespace: -install_requires = - icon4py-common -python_requires = >=3.10 -package_dir = - = src -zip_safe = False - -[options.package_data] -# References: -# https://setuptools.pypa.io/en/latest/userguide/datafiles.html -# https://github.com/abravalheri/experiment-setuptools-package-data -* = *.md, *.rst, *.toml, *.txt, py.typed - -[options.packages.find] -where = src -exclude = - tests diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_global_to_vn.py deleted file mode 100644 index cc6cde483f..0000000000 --- a/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ /dev/null @@ -1,69 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_global_to_vn import ( - apply_nabla2_and_nabla4_global_to_vn, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def apply_nabla2_and_nabla4_global_to_vn_numpy( - area_edge: np.array, - kh_smag_e: np.array, - z_nabla2_e: np.array, - z_nabla4_e2: np.array, - diff_multfac_vn: np.array, - vn: np.array, -): - area_edge = np.expand_dims(area_edge, axis=-1) - diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) - return vn - - -def test_apply_nabla2_and_nabla4_global_to_vn(): - mesh = SimpleMesh() - - area_edge = random_field(mesh, EdgeDim) - kh_smag_e = random_field(mesh, EdgeDim, KDim) - z_nabla2_e = random_field(mesh, EdgeDim, KDim) - z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) - diff_multfac_vn = random_field(mesh, KDim) - vn = random_field(mesh, EdgeDim, KDim) - - vn_ref = apply_nabla2_and_nabla4_global_to_vn_numpy( - np.asarray(area_edge), - np.asarray(kh_smag_e), - np.asarray(z_nabla2_e), - np.asarray(z_nabla4_e2), - np.asarray(diff_multfac_vn), - np.asarray(vn), - ) - - apply_nabla2_and_nabla4_global_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - vn, - offset_provider={}, - ) - assert np.allclose(vn, vn_ref) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_to_vn.py deleted file mode 100644 index 158208caa0..0000000000 --- a/atm_dyn_iconam/tests/test_apply_nabla2_and_nabla4_to_vn.py +++ /dev/null @@ -1,79 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( - apply_nabla2_and_nabla4_to_vn, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def apply_nabla2_and_nabla4_to_vn_numpy( - area_edge: np.array, - kh_smag_e: np.array, - z_nabla2_e: np.array, - z_nabla4_e2: np.array, - diff_multfac_vn: np.array, - nudgecoeff_e: np.array, - vn: np.array, - nudgezone_diff, -): - area_edge = np.expand_dims(area_edge, axis=-1) - diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) - vn = vn + area_edge * ( - np.maximum(nudgezone_diff * nudgecoeff_e, kh_smag_e) * z_nabla2_e - - diff_multfac_vn * z_nabla4_e2 * area_edge - ) - return vn - - -def test_apply_nabla2_and_nabla4_to_vn(): - mesh = SimpleMesh() - - area_edge = random_field(mesh, EdgeDim) - kh_smag_e = random_field(mesh, EdgeDim, KDim) - z_nabla2_e = random_field(mesh, EdgeDim, KDim) - z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) - diff_multfac_vn = random_field(mesh, KDim) - nudgecoeff_e = random_field(mesh, EdgeDim) - vn = random_field(mesh, EdgeDim, KDim) - nudgezone_diff = 9.0 - - vn_ref = apply_nabla2_and_nabla4_to_vn_numpy( - np.asarray(area_edge), - np.asarray(kh_smag_e), - np.asarray(z_nabla2_e), - np.asarray(z_nabla4_e2), - np.asarray(diff_multfac_vn), - np.asarray(nudgecoeff_e), - np.asarray(vn), - nudgezone_diff, - ) - - apply_nabla2_and_nabla4_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - nudgecoeff_e, - vn, - nudgezone_diff, - offset_provider={}, - ) - assert np.allclose(vn, vn_ref) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py deleted file mode 100644 index 3f9e62d5f7..0000000000 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ /dev/null @@ -1,47 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.apply_nabla2_to_vn_in_lateral_boundary import ( - apply_nabla2_to_vn_in_lateral_boundary, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def apply_nabla2_to_vn_in_lateral_boundary_numpy( - z_nabla2_e: np.array, area_edge: np.array, vn: np.array, fac_bdydiff_v -) -> np.array: - area_edge = np.expand_dims(area_edge, axis=-1) - vn = vn + (z_nabla2_e * area_edge * fac_bdydiff_v) - return vn - - -def test_apply_nabla2_to_vn_in_lateral_boundary(): - mesh = SimpleMesh() - - fac_bdydiff_v = 5.0 - z_nabla2_e = random_field(mesh, EdgeDim, KDim) - area_edge = random_field(mesh, EdgeDim) - vn = random_field(mesh, EdgeDim, KDim) - - ref = apply_nabla2_to_vn_in_lateral_boundary_numpy( - np.asarray(z_nabla2_e), np.asarray(area_edge), np.asarray(vn), fac_bdydiff_v - ) - apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v, offset_provider={} - ) - assert np.allclose(vn, ref) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py deleted file mode 100644 index 84227917b2..0000000000 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py +++ /dev/null @@ -1,66 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_apply_nabla2_to_w( - c2e2c0: np.array, - area: np.array, - z_nabla2_c: np.array, - geofac_n2s: np.array, - w: np.array, - diff_multfac_w, -) -> np.array: - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum( - z_nabla2_c[c2e2c0] * geofac_n2s, axis=1 - ) - return w - - -def test_apply_nabla2_to_w(): - mesh = SimpleMesh() - - area = random_field(mesh, CellDim) - z_nabla2_c = random_field(mesh, CellDim, KDim) - geofac_n2s = random_field(mesh, CellDim, C2E2CODim) - w = random_field(mesh, CellDim, KDim) - diff_multfac_w = 5.0 - - ref = mo_apply_nabla2_to_w( - mesh.c2e2cO, - np.asarray(area), - np.asarray(z_nabla2_c), - np.asarray(geofac_n2s), - np.asarray(w), - diff_multfac_w, - ) - - apply_nabla2_to_w( - area, - z_nabla2_c, - geofac_n2s, - w, - diff_multfac_w, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, - ) - - assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py deleted file mode 100644 index 280b77c303..0000000000 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ /dev/null @@ -1,57 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( - apply_nabla2_to_w_in_upper_damping_layer, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def apply_nabla2_to_w_in_upper_damping_layer_numpy( - w: np.array, - diff_multfac_n2w: np.array, - cell_area: np.array, - z_nabla2_c: np.array, -) -> np.array: - cell_area = np.expand_dims(cell_area, axis=-1) - w = w + diff_multfac_n2w * cell_area * z_nabla2_c - return w - - -def test_apply_nabla2_to_w_in_upper_damping_layer(): - mesh = SimpleMesh() - - w = random_field(mesh, CellDim, KDim) - diff_multfac_n2w = random_field(mesh, KDim) - cell_area = random_field(mesh, CellDim) - z_nabla2_c = random_field(mesh, CellDim, KDim) - - ref = apply_nabla2_to_w_in_upper_damping_layer_numpy( - np.asarray(w), - np.asarray(diff_multfac_n2w), - np.asarray(cell_area), - np.asarray(z_nabla2_c), - ) - apply_nabla2_to_w_in_upper_damping_layer( - w, - diff_multfac_n2w, - cell_area, - z_nabla2_c, - offset_provider={}, - ) - assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_calculate_diagnostics_for_turbulence.py b/atm_dyn_iconam/tests/test_calculate_diagnostics_for_turbulence.py deleted file mode 100644 index 0dee48156a..0000000000 --- a/atm_dyn_iconam/tests/test_calculate_diagnostics_for_turbulence.py +++ /dev/null @@ -1,63 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( - calculate_diagnostics_for_turbulence, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def calculate_diagnostics_for_turbulence_numpy( - div: np.array, - k_hc: np.array, - wgtfac_c: np.array, -) -> tuple[np.array, np.array]: - kc_offset_1 = np.roll(k_hc, shift=1, axis=1) - div_offset_1 = np.roll(div, shift=1, axis=1) - div_ic = wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1 - hdef_ic = (wgtfac_c * k_hc + (1.0 - wgtfac_c) * kc_offset_1) ** 2 - return div_ic, hdef_ic - - -def test_calculate_diagnostics_for_turbulence(): - mesh = SimpleMesh() - - wgtfac_c = random_field(mesh, CellDim, KDim) - div = random_field(mesh, CellDim, KDim) - kh_c = random_field(mesh, CellDim, KDim) - - div_ic = zero_field(mesh, CellDim, KDim) - hdef_ic = zero_field(mesh, CellDim, KDim) - - div_ref, kh_c_ref = calculate_diagnostics_for_turbulence_numpy( - np.asarray(div), - np.asarray(kh_c), - np.asarray(wgtfac_c), - ) - - calculate_diagnostics_for_turbulence( - div, - kh_c, - wgtfac_c, - div_ic, - hdef_ic, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(hdef_ic[:, 1:], kh_c_ref[:, 1:]) - assert np.allclose(div_ic[:, 1:], div_ref[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py deleted file mode 100644 index 0909deb3b5..0000000000 --- a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ /dev/null @@ -1,61 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( - calculate_horizontal_gradients_for_turbulence, -) -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def calculate_horizontal_gradients_for_turbulence_numpy( - c2e2c0: np.array, - w: np.array, - geofac_grg_x: np.array, - geofac_grg_y: np.array, -) -> tuple[np.array]: - geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) - dwdx = np.sum(geofac_grg_x * w[c2e2c0], axis=1) - - geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) - dwdy = np.sum(geofac_grg_y * w[c2e2c0], axis=1) - return dwdx, dwdy - - -def test_calculate_horizontal_gradients_for_turbulence(): - mesh = SimpleMesh() - - w = random_field(mesh, CellDim, KDim) - geofac_grg_x = random_field(mesh, CellDim, C2E2CODim) - geofac_grg_y = random_field(mesh, CellDim, C2E2CODim) - dwdx = zero_field(mesh, CellDim, KDim) - dwdy = zero_field(mesh, CellDim, KDim) - - dwdx_ref, dwdy_ref = calculate_horizontal_gradients_for_turbulence_numpy( - mesh.c2e2cO, np.asarray(w), np.asarray(geofac_grg_x), np.asarray(geofac_grg_y) - ) - calculate_horizontal_gradients_for_turbulence( - w, - geofac_grg_x, - geofac_grg_y, - dwdx, - dwdy, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, - ) - - assert np.allclose(dwdx, dwdx_ref) - assert np.allclose(dwdy, dwdy_ref) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py deleted file mode 100644 index 9de6f08837..0000000000 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py +++ /dev/null @@ -1,47 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def calculate_nabla2_for_w_numpy( - c2e2cO: np.array, w: np.array, geofac_n2s: np.array -) -> np.array: - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - z_nabla2_c = np.sum(w[c2e2cO] * geofac_n2s, axis=1) - return z_nabla2_c - - -def test_calculate_nabla2_for_w(): - mesh = SimpleMesh() - - w = random_field(mesh, CellDim, KDim) - geofac_n2s = random_field(mesh, CellDim, C2E2CODim) - z_nabla2_c = zero_field(mesh, CellDim, KDim) - - ref = calculate_nabla2_for_w_numpy( - mesh.c2e2cO, np.asarray(w), np.asarray(geofac_n2s) - ) - calculate_nabla2_for_w( - w, - geofac_n2s, - z_nabla2_c, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, - ) - assert np.allclose(z_nabla2_c, ref) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_for_z.py b/atm_dyn_iconam/tests/test_calculate_nabla2_for_z.py deleted file mode 100644 index 10229e638f..0000000000 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_for_z.py +++ /dev/null @@ -1,62 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import calculate_nabla2_for_z -from icon4py.common.dimension import CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def calculate_nabla2_for_z_numpy( - e2c: np.array, - kh_smag_e: np.array, - inv_dual_edge_length: np.array, - theta_v: np.array, -) -> np.array: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - - theta_v_e2c = theta_v[e2c] - theta_v_weighted = theta_v_e2c[:, 1] - theta_v_e2c[:, 0] - - z_nabla2_e = kh_smag_e * inv_dual_edge_length * theta_v_weighted - return z_nabla2_e - - -def test_calculate_nabla2_for_z(): - mesh = SimpleMesh() - - kh_smag_e = random_field(mesh, EdgeDim, KDim) - inv_dual_edge_length = random_field(mesh, EdgeDim) - theta_v = random_field(mesh, CellDim, KDim) - z_nabla2_e = random_field(mesh, EdgeDim, KDim) - - ref = calculate_nabla2_for_z_numpy( - mesh.e2c, - np.asarray(kh_smag_e), - np.asarray(inv_dual_edge_length), - np.asarray(theta_v), - ) - - calculate_nabla2_for_z( - kh_smag_e, - inv_dual_edge_length, - theta_v, - z_nabla2_e, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - }, - ) - assert np.allclose(z_nabla2_e, ref) diff --git a/atm_dyn_iconam/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/atm_dyn_iconam/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py deleted file mode 100644 index 30cfd73d1d..0000000000 --- a/atm_dyn_iconam/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ /dev/null @@ -1,47 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( - enhance_diffusion_coefficient_for_grid_point_cold_pools, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def enhance_diffusion_coefficient_for_grid_point_cold_pools_numpy( - e2c: np.array, - kh_smag_e: np.array, - enh_diffu_3d: np.array, -) -> np.array: - kh_smag_e = np.maximum(kh_smag_e, np.max(enh_diffu_3d[e2c], axis=1)) - return kh_smag_e - - -def test_enhance_diffusion_coefficient_for_grid_point_cold_pools(): - mesh = SimpleMesh() - kh_smag_e = random_field(mesh, EdgeDim, KDim) - enh_diffu_3d = random_field(mesh, CellDim, KDim) - - kh_smag_e_ref = enhance_diffusion_coefficient_for_grid_point_cold_pools_numpy( - mesh.e2c, np.asarray(kh_smag_e), np.asarray(enh_diffu_3d) - ) - - enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d, offset_provider={"E2C": mesh.get_e2c_offset_provider()} - ) - - np.allclose(kh_smag_e_ref, kh_smag_e) diff --git a/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py deleted file mode 100644 index 9676e3c2e8..0000000000 --- a/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ /dev/null @@ -1,50 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, -) -from icon4py.common.dimension import CellDim, KDim, V2CDim, VertexDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - v2c: np.array, p_cell_in: np.array, c_intp: np.array -) -> np.array: - c_intp = np.expand_dims(c_intp, axis=-1) - p_vert_out = np.sum(p_cell_in[v2c] * c_intp, axis=1) - return p_vert_out - - -def test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(): - mesh = SimpleMesh() - - p_cell_in = random_field(mesh, CellDim, KDim) - c_intp = random_field(mesh, VertexDim, V2CDim) - p_vert_out = zero_field(mesh, VertexDim, KDim) - - ref = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - mesh.v2c, np.asarray(p_cell_in), np.asarray(c_intp) - ) - mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in, - c_intp, - p_vert_out, - offset_provider={"V2C": mesh.get_v2c_offset_provider()}, - ) - - assert np.allclose(p_vert_out, ref) diff --git a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py deleted file mode 100644 index cf621cd1a4..0000000000 --- a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ /dev/null @@ -1,59 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.common.dimension import EdgeDim, KDim, V2EDim, VertexDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_intp_rbf_rbf_vec_interpol_vertex_numpy( - v2e: np.array, p_e_in: np.array, ptr_coeff_1: np.array, ptr_coeff_2: np.array -) -> tuple[np.array]: - ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) - p_v_out = np.sum(p_e_in[v2e] * ptr_coeff_1, axis=1) - - ptr_coeff_2 = np.expand_dims(ptr_coeff_2, axis=-1) - p_u_out = np.sum(p_e_in[v2e] * ptr_coeff_2, axis=1) - - return p_v_out, p_u_out - - -def test_mo_intp_rbf_rbf_vec_interpol_vertex(): - mesh = SimpleMesh() - - p_e_in = random_field(mesh, EdgeDim, KDim) - ptr_coeff_1 = random_field(mesh, VertexDim, V2EDim) - ptr_coeff_2 = random_field(mesh, VertexDim, V2EDim) - p_v_out = zero_field(mesh, VertexDim, KDim) - p_u_out = zero_field(mesh, VertexDim, KDim) - - p_v_out_ref, p_u_out_ref = mo_intp_rbf_rbf_vec_interpol_vertex_numpy( - mesh.v2e, np.asarray(p_e_in), np.asarray(ptr_coeff_1), np.asarray(ptr_coeff_2) - ) - mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, - ptr_coeff_1, - ptr_coeff_2, - p_v_out, - p_u_out, - offset_provider={"V2E": mesh.get_v2e_offset_provider()}, - ) - - assert np.allclose(p_v_out, p_v_out_ref) - assert np.allclose(p_u_out, p_u_out_ref) diff --git a/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py deleted file mode 100644 index 07deadda10..0000000000 --- a/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ /dev/null @@ -1,55 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_math_divrot_rot_vertex_ri_dsl import ( - mo_math_divrot_rot_vertex_ri_dsl, -) -from icon4py.common.dimension import EdgeDim, KDim, V2EDim, VertexDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_math_divrot_rot_vertex_ri_dsl_numpy( - v2e: np.array, vec_e: np.array, geofac_rot: np.array -) -> np.array: - geofac_rot = np.expand_dims(geofac_rot, axis=-1) - rot_vec = np.sum(vec_e[v2e] * geofac_rot, axis=1) - return rot_vec - - -def test_mo_math_divrot_rot_vertex_ri_dsl_numpy(): - mesh = SimpleMesh() - - vec_e = random_field(mesh, EdgeDim, KDim) - geofac_rot = random_field(mesh, VertexDim, V2EDim) - rot_vec = zero_field(mesh, VertexDim, KDim) - - ref = mo_math_divrot_rot_vertex_ri_dsl_numpy( - mesh.v2e, - np.asarray(vec_e), - np.asarray(geofac_rot), - ) - - mo_math_divrot_rot_vertex_ri_dsl( - vec_e, - geofac_rot, - rot_vec, - offset_provider={ - "V2E": mesh.get_v2e_offset_provider(), - }, - ) - - assert np.allclose(rot_vec, ref) diff --git a/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py b/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py deleted file mode 100644 index 40227d5568..0000000000 --- a/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py +++ /dev/null @@ -1,83 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_math_gradients_grad_green_gauss_cell_dsl import ( - mo_math_gradients_grad_green_gauss_cell_dsl, -) -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_math_gradients_grad_green_gauss_cell_dsl_numpy( - c2e2cO: np.array, - p_ccpr1: np.array, - p_ccpr2: np.array, - geofac_grg_x: np.array, - geofac_grg_y: np.array, -) -> tuple[np.array]: - geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) - p_grad_1_u = np.sum(geofac_grg_x * p_ccpr1[c2e2cO], axis=1) - - geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) - p_grad_1_v = np.sum(geofac_grg_y * p_ccpr1[c2e2cO], axis=1) - - p_grad_2_u = np.sum(geofac_grg_x * p_ccpr2[c2e2cO], axis=1) - - p_grad_2_v = np.sum(geofac_grg_y * p_ccpr2[c2e2cO], axis=1) - return p_grad_1_u, p_grad_1_v, p_grad_2_u, p_grad_2_v - - -def test_mo_math_gradients_grad_green_gauss_cell_dsl_numpy(): - mesh = SimpleMesh() - - p_ccpr1 = random_field(mesh, CellDim, KDim) - p_ccpr2 = random_field(mesh, CellDim, KDim) - geofac_grg_x = random_field(mesh, CellDim, C2E2CODim) - geofac_grg_y = random_field(mesh, CellDim, C2E2CODim) - p_grad_1_u = zero_field(mesh, CellDim, KDim) - p_grad_1_v = zero_field(mesh, CellDim, KDim) - p_grad_2_u = zero_field(mesh, CellDim, KDim) - p_grad_2_v = zero_field(mesh, CellDim, KDim) - - ( - p_grad_1_u_ref, - p_grad_1_v_ref, - p_grad_2_u_ref, - p_grad_2_v_ref, - ) = mo_math_gradients_grad_green_gauss_cell_dsl_numpy( - mesh.c2e2cO, - np.asarray(p_ccpr1), - np.asarray(p_ccpr2), - np.asarray(geofac_grg_x), - np.asarray(geofac_grg_y), - ) - - mo_math_gradients_grad_green_gauss_cell_dsl( - p_grad_1_u, - p_grad_1_v, - p_grad_2_u, - p_grad_2_v, - p_ccpr1, - p_ccpr2, - geofac_grg_x, - geofac_grg_y, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, - ) - assert np.allclose(p_grad_1_u_ref, p_grad_1_u) - assert np.allclose(p_grad_1_v_ref, p_grad_1_v) - assert np.allclose(p_grad_2_u_ref, p_grad_2_u) - assert np.allclose(p_grad_2_v_ref, p_grad_2_v) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_4th_order_divdamp.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_4th_order_divdamp.py deleted file mode 100644 index a3718d01b3..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_4th_order_divdamp.py +++ /dev/null @@ -1,53 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_4th_order_divdamp import ( - mo_solve_nonhydro_4th_order_divdamp, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_4th_order_divdamp_numpy( - scal_divdamp: np.array, - z_graddiv2_vn: np.array, - vn: np.array, -) -> np.array: - scal_divdamp = np.expand_dims(scal_divdamp, axis=0) - vn = vn + (scal_divdamp * z_graddiv2_vn) - return vn - - -def test_mo_solve_nonhydro_4th_order_divdamp(): - mesh = SimpleMesh() - - scal_divdamp = random_field(mesh, KDim) - z_graddiv2_vn = random_field(mesh, EdgeDim, KDim) - vn = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_4th_order_divdamp_numpy( - np.asarray(scal_divdamp), - np.asarray(z_graddiv2_vn), - np.asarray(vn), - ) - mo_solve_nonhydro_4th_order_divdamp( - scal_divdamp, - z_graddiv2_vn, - vn, - offset_provider={}, - ) - assert np.allclose(vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_01.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_01.py deleted file mode 100644 index cc4f45e757..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_01.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_01 import ( - mo_solve_nonhydro_stencil_01, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_01_numpy( - z_rth_pr_1: np.array, - z_rth_pr_2: np.array, -) -> tuple[np.array]: - z_rth_pr_1 = np.zeros_like(z_rth_pr_1) - z_rth_pr_2 = np.zeros_like(z_rth_pr_2) - return z_rth_pr_1, z_rth_pr_2 - - -def test_mo_solve_nonhydro_stencil_01_z_rth_pr_2(): - mesh = SimpleMesh() - - z_rth_pr_1 = random_field(mesh, CellDim, KDim) - z_rth_pr_2 = random_field(mesh, CellDim, KDim) - - z_rth_pr_1_ref, z_rth_pr_2_ref = mo_solve_nonhydro_stencil_01_numpy( - z_rth_pr_1, z_rth_pr_2 - ) - mo_solve_nonhydro_stencil_01( - z_rth_pr_1, - z_rth_pr_2, - offset_provider={}, - ) - assert np.allclose(z_rth_pr_1, z_rth_pr_1_ref) - assert np.allclose(z_rth_pr_2, z_rth_pr_2_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_02.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_02.py deleted file mode 100644 index df51459407..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_02.py +++ /dev/null @@ -1,52 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_02 import ( - mo_solve_nonhydro_stencil_02, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_02_numpy( - exner: np.array, exner_ref_mc: np.array, exner_pr: np.array, exner_exfac: np.array -) -> tuple[np.array]: - z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr - exner_pr = exner - exner_ref_mc - return z_exner_ex_pr, exner_pr - - -def test_mo_solve_nonhydro_stencil_02(): - mesh = SimpleMesh() - - exner = random_field(mesh, CellDim, KDim) - exner_ref_mc = random_field(mesh, CellDim, KDim) - exner_pr = zero_field(mesh, CellDim, KDim) - exner_exfac = random_field(mesh, CellDim, KDim) - z_exner_ex_pr = zero_field(mesh, CellDim, KDim) - - z_exner_ex_pr_ref, exner_pr_ref = mo_solve_nonhydro_stencil_02_numpy( - np.asarray(exner), - np.asarray(exner_ref_mc), - np.asarray(exner_pr), - np.asarray(exner_exfac), - ) - mo_solve_nonhydro_stencil_02( - exner_exfac, exner, exner_ref_mc, exner_pr, z_exner_ex_pr, offset_provider={} - ) - assert np.allclose(z_exner_ex_pr_ref, z_exner_ex_pr) - assert np.allclose(exner_pr_ref, exner_pr) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_03.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_03.py deleted file mode 100644 index 69c3feefb7..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_03.py +++ /dev/null @@ -1,44 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_03 import ( - mo_solve_nonhydro_stencil_03, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_03_numpy( - z_exner_ex_pr: np.array, -) -> np.array: - z_exner_ex_pr = np.zeros_like(z_exner_ex_pr) - return z_exner_ex_pr - - -def test_mo_solve_nonhydro_stencil_03(): - mesh = SimpleMesh() - z_exner_ex_pr = random_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_03_numpy( - np.asarray(z_exner_ex_pr), - ) - - mo_solve_nonhydro_stencil_03( - z_exner_ex_pr, - offset_provider={}, - ) - assert np.allclose(z_exner_ex_pr, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_04.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_04.py deleted file mode 100644 index 244f9bdca7..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_04.py +++ /dev/null @@ -1,55 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_04 import ( - mo_solve_nonhydro_stencil_04, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_04_numpy( - wgtfacq_c: np.array, - z_exner_ex_pr: np.array, -) -> np.array: - z_exner_ic = ( - np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) - + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) - + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) - ) - return z_exner_ic - - -def test_mo_solve_nonhydro_stencil_04(): - mesh = SimpleMesh() - z_exner_ex_pr = random_field(mesh, CellDim, KDim) - wgtfacq_c = random_field(mesh, CellDim, KDim) - z_exner_ic = zero_field(mesh, CellDim, KDim) - - z_exner_ic_ref = mo_solve_nonhydro_stencil_04_numpy( - np.asarray(wgtfacq_c), - np.asarray(z_exner_ex_pr), - ) - - mo_solve_nonhydro_stencil_04( - z_exner_ex_pr, - wgtfacq_c, - z_exner_ic, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_exner_ic, z_exner_ic_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_06.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_06.py deleted file mode 100644 index b0c74ff8e6..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_06.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_06 import ( - mo_solve_nonhydro_stencil_06, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_06_numpy( - z_exner_ic: np.array, - inv_ddqz_z_full: np.array, -) -> np.array: - z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full - return z_dexner_dz_c_1 - - -def test_mo_solve_nonhydro_stencil_06(): - mesh = SimpleMesh() - z_exner_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - z_dexner_dz_c_1 = zero_field(mesh, CellDim, KDim) - - z_dexner_dz_c_1_ref = mo_solve_nonhydro_stencil_06_numpy( - np.asarray(z_exner_ic), - np.asarray(inv_ddqz_z_full), - ) - - mo_solve_nonhydro_stencil_06( - z_exner_ic, - inv_ddqz_z_full, - z_dexner_dz_c_1, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_dexner_dz_c_1, z_dexner_dz_c_1_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_07.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_07.py deleted file mode 100644 index 5e83ad278a..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_07.py +++ /dev/null @@ -1,60 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_07 import ( - mo_solve_nonhydro_stencil_07, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_07_numpy( - rho: np.array, rho_ref_mc: np.array, theta_v: np.array, theta_ref_mc: np.array -) -> tuple[np.array]: - z_rth_pr_1 = rho - rho_ref_mc - z_rth_pr_2 = theta_v - theta_ref_mc - return z_rth_pr_1, z_rth_pr_2 - - -def test_mo_solve_nonhydro_stencil_07(): - mesh = SimpleMesh() - - rho = random_field(mesh, CellDim, KDim) - rho_ref_mc = random_field(mesh, CellDim, KDim) - theta_v = random_field(mesh, CellDim, KDim) - theta_ref_mc = random_field(mesh, CellDim, KDim) - z_rth_pr_1 = zero_field(mesh, CellDim, KDim) - z_rth_pr_2 = zero_field(mesh, CellDim, KDim) - - z_rth_pr_1_ref, z_rth_pr_2_ref = mo_solve_nonhydro_stencil_07_numpy( - np.asarray(rho), - np.asarray(rho_ref_mc), - np.asarray(theta_v), - np.asarray(theta_ref_mc), - ) - mo_solve_nonhydro_stencil_07( - rho, - rho_ref_mc, - theta_v, - theta_ref_mc, - z_rth_pr_1, - z_rth_pr_2, - offset_provider={}, - ) - - assert np.allclose(z_rth_pr_1, z_rth_pr_1_ref) - assert np.allclose(z_rth_pr_2, z_rth_pr_2_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_lower.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_lower.py deleted file mode 100644 index 7177ebfafb..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_lower.py +++ /dev/null @@ -1,43 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_11_lower import ( - mo_solve_nonhydro_stencil_11_lower, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_11_lower_numpy() -> np.array: - z_theta_v_pr_ic = 0 - - return z_theta_v_pr_ic - - -def test_mo_solve_nonhydro_stencil_11_lower(): - mesh = SimpleMesh() - - z_theta_v_pr_ic = random_field(mesh, CellDim, KDim) - - z_theta_v_pr_ic_ref = mo_solve_nonhydro_stencil_11_lower_numpy() - - mo_solve_nonhydro_stencil_11_lower( - z_theta_v_pr_ic, - offset_provider={}, - ) - - assert np.allclose(z_theta_v_pr_ic, z_theta_v_pr_ic_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_12.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_12.py deleted file mode 100644 index afb1a85839..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_12.py +++ /dev/null @@ -1,65 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_12 import ( - mo_solve_nonhydro_stencil_12, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_12_numpy( - z_theta_v_pr_ic: np.array, - d2dexdz2_fac1_mc: np.array, - d2dexdz2_fac2_mc: np.array, - z_rth_pr_2: np.array, -) -> np.array: - z_theta_v_pr_ic_offset_1 = z_theta_v_pr_ic[:, 1:] - z_dexner_dz_c_2 = -0.5 * ( - (z_theta_v_pr_ic[:, :-1] - z_theta_v_pr_ic_offset_1) * d2dexdz2_fac1_mc - + z_rth_pr_2 * d2dexdz2_fac2_mc - ) - return z_dexner_dz_c_2 - - -def test_mo_solve_nonhydro_stencil_12(): - mesh = SimpleMesh() - - z_theta_v_pr_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - d2dexdz2_fac1_mc = random_field(mesh, CellDim, KDim) - z_rth_pr_2 = random_field(mesh, CellDim, KDim) - d2dexdz2_fac2_mc = random_field(mesh, CellDim, KDim) - - z_dexner_dz_c_2 = zero_field(mesh, CellDim, KDim) - - z_dexner_dz_c_2_ref = mo_solve_nonhydro_stencil_12_numpy( - np.asarray(z_theta_v_pr_ic), - np.asarray(d2dexdz2_fac1_mc), - np.asarray(d2dexdz2_fac2_mc), - np.asarray(z_rth_pr_2), - ) - - mo_solve_nonhydro_stencil_12( - z_theta_v_pr_ic, - d2dexdz2_fac1_mc, - d2dexdz2_fac2_mc, - z_rth_pr_2, - z_dexner_dz_c_2, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_dexner_dz_c_2, z_dexner_dz_c_2_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_13.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_13.py deleted file mode 100644 index 86fec6b518..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_13.py +++ /dev/null @@ -1,60 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_13 import ( - mo_solve_nonhydro_stencil_13, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_13_numpy( - rho: np.array, rho_ref_mc: np.array, theta_v: np.array, theta_ref_mc: np.array -) -> tuple[np.array]: - z_rth_pr_1 = rho - rho_ref_mc - z_rth_pr_2 = theta_v - theta_ref_mc - return z_rth_pr_1, z_rth_pr_2 - - -def test_mo_solve_nonhydro_stencil_13(): - mesh = SimpleMesh() - - rho = random_field(mesh, CellDim, KDim) - rho_ref_mc = random_field(mesh, CellDim, KDim) - theta_v = random_field(mesh, CellDim, KDim) - theta_ref_mc = random_field(mesh, CellDim, KDim) - z_rth_pr_1 = zero_field(mesh, CellDim, KDim) - z_rth_pr_2 = zero_field(mesh, CellDim, KDim) - - z_rth_pr_1_ref, z_rth_pr_2_ref = mo_solve_nonhydro_stencil_13_numpy( - np.asarray(rho), - np.asarray(rho_ref_mc), - np.asarray(theta_v), - np.asarray(theta_ref_mc), - ) - mo_solve_nonhydro_stencil_13( - rho, - rho_ref_mc, - theta_v, - theta_ref_mc, - z_rth_pr_1, - z_rth_pr_2, - offset_provider={}, - ) - - assert np.allclose(z_rth_pr_1, z_rth_pr_1_ref) - assert np.allclose(z_rth_pr_2, z_rth_pr_2_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_14.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_14.py deleted file mode 100644 index 39f926ddb8..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_14.py +++ /dev/null @@ -1,47 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_14 import ( - mo_solve_nonhydro_stencil_14, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_14_numpy( - z_rho_e: np.array, z_theta_v_e: np.array -) -> tuple[np.array]: - z_rho_e = np.zeros_like(z_rho_e) - z_theta_v_e = np.zeros_like(z_theta_v_e) - return z_rho_e, z_theta_v_e - - -def test_mo_solve_nonhydro_stencil_14_z_theta_v_e(): - mesh = SimpleMesh() - - z_rho_e = zero_field(mesh, EdgeDim, KDim) - z_theta_v_e = zero_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_14_numpy( - np.asarray(z_rho_e), np.asarray(z_theta_v_e) - ) - mo_solve_nonhydro_stencil_14( - z_rho_e, - z_theta_v_e, - offset_provider={}, - ) - assert np.allclose(z_theta_v_e, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_15.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_15.py deleted file mode 100644 index 16c1ac5fb1..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_15.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_15 import ( - mo_solve_nonhydro_stencil_15, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_15_numpy( - z_rho_e: np.array, z_theta_v_e: np.array -) -> tuple[np.array]: - z_rho_e = np.zeros_like(z_rho_e) - z_theta_v_e = np.zeros_like(z_theta_v_e) - return z_rho_e, z_theta_v_e - - -def test_mo_solve_nonhydro_stencil_15_z_theta_v_e(): - mesh = SimpleMesh() - - z_rho_e = random_field(mesh, EdgeDim, KDim) - z_theta_v_e = random_field(mesh, EdgeDim, KDim) - - z_rho_e_ref, z_theta_v_e_ref = mo_solve_nonhydro_stencil_15_numpy( - np.asarray(z_rho_e), np.asarray(z_theta_v_e) - ) - - mo_solve_nonhydro_stencil_15( - z_rho_e, - z_theta_v_e, - offset_provider={}, - ) - assert np.allclose(z_rho_e, z_rho_e_ref) - assert np.allclose(z_theta_v_e, z_theta_v_e_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_17.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_17.py deleted file mode 100644 index c55e276b56..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_17.py +++ /dev/null @@ -1,74 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_17 import ( - mo_solve_nonhydro_stencil_17, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_17_numpy( - e2c: np.array, - hmask_dd3d: np.array, - scalfac_dd3d: np.array, - inv_dual_edge_length: np.array, - z_dwdz_dd: np.array, - z_graddiv_vn: np.array, -) -> np.array: - scalfac_dd3d = np.expand_dims(scalfac_dd3d, axis=0) - hmask_dd3d = np.expand_dims(hmask_dd3d, axis=-1) - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - - z_dwdz_dd_e2c = z_dwdz_dd[e2c] - z_dwdz_dd_weighted = z_dwdz_dd_e2c[:, 1] - z_dwdz_dd_e2c[:, 0] - - z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * z_dwdz_dd_weighted - ) - return z_graddiv_vn - - -def test_mo_solve_nonhydro_stencil_17(): - mesh = SimpleMesh() - - hmask_dd3d = random_field(mesh, EdgeDim) - scalfac_dd3d = random_field(mesh, KDim) - inv_dual_edge_length = random_field(mesh, EdgeDim) - z_dwdz_dd = random_field(mesh, CellDim, KDim) - z_graddiv_vn = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_17_numpy( - mesh.e2c, - np.asarray(hmask_dd3d), - np.asarray(scalfac_dd3d), - np.asarray(inv_dual_edge_length), - np.asarray(z_dwdz_dd), - np.asarray(z_graddiv_vn), - ) - - mo_solve_nonhydro_stencil_17( - hmask_dd3d, - scalfac_dd3d, - inv_dual_edge_length, - z_dwdz_dd, - z_graddiv_vn, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - }, - ) - assert np.allclose(z_graddiv_vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_18.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_18.py deleted file mode 100644 index 8555bab3ed..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_18.py +++ /dev/null @@ -1,60 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_18 import ( - mo_solve_nonhydro_stencil_18, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_18_numpy( - e2c: np.array, - inv_dual_edge_length: np.array, - z_exner_ex_pr: np.array, -) -> np.array: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - - z_exner_ex_pr_e2c = z_exner_ex_pr[e2c] - z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - - z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - return z_gradh_exner - - -def test_mo_solve_nonhydro_stencil_18(): - mesh = SimpleMesh() - - inv_dual_edge_length = random_field(mesh, EdgeDim) - z_exner_ex_pr = random_field(mesh, CellDim, KDim) - z_gradh_exner = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_18_numpy( - mesh.e2c, - np.asarray(inv_dual_edge_length), - np.asarray(z_exner_ex_pr), - ) - - mo_solve_nonhydro_stencil_18( - inv_dual_edge_length, - z_exner_ex_pr, - z_gradh_exner, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - }, - ) - assert np.allclose(z_gradh_exner, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py deleted file mode 100644 index cd163dc2e2..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py +++ /dev/null @@ -1,75 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_19 import ( - mo_solve_nonhydro_stencil_19, -) -from icon4py.common.dimension import CellDim, E2CDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_19_numpy( - e2c: np.array, - inv_dual_edge_length: np.array, - z_exner_ex_pr: np.array, - ddxn_z_full: np.array, - c_lin_e: np.array, - z_dexner_dz_c_1: np.array, -) -> np.array: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - - z_exner_ex_pr_e2c = z_exner_ex_pr[e2c] - z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - - z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( - c_lin_e * z_dexner_dz_c_1[e2c], axis=1 - ) - return z_gradh_exner - - -def test_mo_solve_nonhydro_stencil_19(): - mesh = SimpleMesh() - - inv_dual_edge_length = random_field(mesh, EdgeDim) - z_exner_ex_pr = random_field(mesh, CellDim, KDim) - ddxn_z_full = random_field(mesh, EdgeDim, KDim) - c_lin_e = random_field(mesh, EdgeDim, E2CDim) - z_dexner_dz_c_1 = random_field(mesh, CellDim, KDim) - z_gradh_exner = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_19_numpy( - mesh.e2c, - np.asarray(inv_dual_edge_length), - np.asarray(z_exner_ex_pr), - np.asarray(ddxn_z_full), - np.asarray(c_lin_e), - np.asarray(z_dexner_dz_c_1), - ) - - mo_solve_nonhydro_stencil_19( - inv_dual_edge_length, - z_exner_ex_pr, - ddxn_z_full, - c_lin_e, - z_dexner_dz_c_1, - z_gradh_exner, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - }, - ) - assert np.allclose(z_gradh_exner, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_22.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_22.py deleted file mode 100644 index 41bc0305c4..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_22.py +++ /dev/null @@ -1,59 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_22 import ( - mo_solve_nonhydro_stencil_22, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, random_mask -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_22_numpy( - ipeidx_dsl: np.array, - pg_exdist: np.array, - z_hydro_corr: np.array, - z_gradh_exner: np.array, -) -> np.array: - z_hydro_corr = np.expand_dims(z_hydro_corr, axis=-1) - z_gradh_exner = np.where( - ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner - ) - return z_gradh_exner - - -def test_mo_solve_nonhydro_stencil_22(): - mesh = SimpleMesh() - - ipeidx_dsl = random_mask(mesh, EdgeDim, KDim) - pg_exdist = random_field(mesh, EdgeDim, KDim) - z_hydro_corr = random_field(mesh, EdgeDim) - z_gradh_exner = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_22_numpy( - np.asarray(ipeidx_dsl), - np.asarray(pg_exdist), - np.asarray(z_hydro_corr), - np.asarray(z_gradh_exner), - ) - mo_solve_nonhydro_stencil_22( - ipeidx_dsl, - pg_exdist, - z_hydro_corr, - z_gradh_exner, - offset_provider={}, - ) - assert np.allclose(z_gradh_exner, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_23.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_23.py deleted file mode 100644 index bf367eeed8..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_23.py +++ /dev/null @@ -1,87 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_23 import ( - mo_solve_nonhydro_stencil_23, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_23_numpy( - vn_nnow: np.array, - ddt_vn_adv_ntl1: np.array, - ddt_vn_adv_ntl2: np.array, - ddt_vn_phy: np.array, - z_theta_v_e: np.array, - z_gradh_exner: np.array, - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, -) -> np.array: - vn_nnew = vn_nnow + dtime * ( - wgt_nnow_vel * ddt_vn_adv_ntl1 - + wgt_nnew_vel * ddt_vn_adv_ntl2 - + ddt_vn_phy - - cpd * z_theta_v_e * z_gradh_exner - ) - return vn_nnew - - -def test_mo_solve_nonhydro_stencil_23(): - mesh = SimpleMesh() - - vn_nnow = random_field(mesh, EdgeDim, KDim) - ddt_vn_adv_ntl1 = random_field(mesh, EdgeDim, KDim) - ddt_vn_adv_ntl2 = random_field(mesh, EdgeDim, KDim) - ddt_vn_phy = random_field(mesh, EdgeDim, KDim) - z_theta_v_e = random_field(mesh, EdgeDim, KDim) - z_gradh_exner = random_field(mesh, EdgeDim, KDim) - vn_nnew = zero_field(mesh, EdgeDim, KDim) - dtime = 5.0 - wgt_nnow_vel = 8.0 - wgt_nnew_vel = 7.0 - cpd = 2.0 - - ref = mo_solve_nonhydro_stencil_23_numpy( - np.asarray(vn_nnow), - np.asarray(ddt_vn_adv_ntl1), - np.asarray(ddt_vn_adv_ntl2), - np.asarray(ddt_vn_phy), - np.asarray(z_theta_v_e), - np.asarray(z_gradh_exner), - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, - ) - mo_solve_nonhydro_stencil_23( - vn_nnow, - ddt_vn_adv_ntl1, - ddt_vn_adv_ntl2, - ddt_vn_phy, - z_theta_v_e, - z_gradh_exner, - vn_nnew, - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, - offset_provider={}, - ) - assert np.allclose(vn_nnew, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_24.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_24.py deleted file mode 100644 index 7a89e37ed0..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_24.py +++ /dev/null @@ -1,71 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_24 import ( - mo_solve_nonhydro_stencil_24, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_24_numpy( - vn_nnow: np.array, - ddt_vn_adv_ntl1: np.array, - ddt_vn_phy: np.array, - z_theta_v_e: np.array, - z_gradh_exner: np.array, - dtime: float, - cpd: float, -) -> np.array: - vn_nnew = vn_nnow + dtime * ( - ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner - ) - return vn_nnew - - -def test_mo_solve_nonhydro_stencil_24(): - mesh = SimpleMesh() - - dtime, cpd = 10.0, 10.0 - vn_nnow = random_field(mesh, EdgeDim, KDim) - ddt_vn_adv_ntl1 = random_field(mesh, EdgeDim, KDim) - ddt_vn_phy = random_field(mesh, EdgeDim, KDim) - z_theta_v_e = random_field(mesh, EdgeDim, KDim) - z_gradh_exner = random_field(mesh, EdgeDim, KDim) - vn_nnew = zero_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_24_numpy( - np.asarray(vn_nnow), - np.asarray(ddt_vn_adv_ntl1), - np.asarray(ddt_vn_phy), - np.asarray(z_theta_v_e), - np.asarray(z_gradh_exner), - dtime, - cpd, - ) - mo_solve_nonhydro_stencil_24( - vn_nnow, - ddt_vn_adv_ntl1, - ddt_vn_phy, - z_theta_v_e, - z_gradh_exner, - vn_nnew, - dtime, - cpd, - offset_provider={}, - ) - assert np.allclose(vn_nnew, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py deleted file mode 100644 index 56120f26e0..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_25 import ( - mo_solve_nonhydro_stencil_25, -) -from icon4py.common.dimension import E2C2EODim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_25_numpy( - e2c2eO: np.array, geofac_grdiv: np.array, z_graddiv_vn: np.array -) -> np.array: - geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) - z_temp = np.sum(z_graddiv_vn[e2c2eO] * geofac_grdiv, axis=1) - return z_temp - - -def test_mo_solve_nonhydro_stencil_25(): - mesh = SimpleMesh() - - z_graddiv_vn = random_field(mesh, EdgeDim, KDim) - geofac_grdiv = random_field(mesh, EdgeDim, E2C2EODim) - z_graddiv2_vn = zero_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_25_numpy( - mesh.e2c2eO, np.asarray(geofac_grdiv), np.asarray(z_graddiv_vn) - ) - mo_solve_nonhydro_stencil_25( - geofac_grdiv, - z_graddiv_vn, - z_graddiv2_vn, - offset_provider={"E2C2EO": mesh.get_e2c2eO_offset_provider()}, - ) - assert np.allclose(z_graddiv2_vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_26.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_26.py deleted file mode 100644 index 772ef4bac8..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_26.py +++ /dev/null @@ -1,48 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_26 import ( - mo_solve_nonhydro_stencil_26, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_26_numpy( - z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2 -) -> np.array: - vn = vn + (scal_divdamp_o2 * z_graddiv_vn) - return vn - - -def test_mo_solve_nonhydro_stencil_26(): - mesh = SimpleMesh() - - z_graddiv_vn = random_field(mesh, EdgeDim, KDim) - vn = random_field(mesh, EdgeDim, KDim) - scal_divdamp_o2 = 5.0 - - ref = mo_solve_nonhydro_stencil_26_numpy( - np.asarray(z_graddiv_vn), np.asarray(vn), scal_divdamp_o2 - ) - mo_solve_nonhydro_stencil_26( - z_graddiv_vn, - vn, - scal_divdamp_o2, - offset_provider={}, - ) - assert np.allclose(vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_27.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_27.py deleted file mode 100644 index 19f91c958a..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_27.py +++ /dev/null @@ -1,61 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_27 import ( - mo_solve_nonhydro_stencil_27, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_27_numpy( - scal_divdamp: np.array, - bdy_divdamp: np.array, - nudgecoeff_e: np.array, - z_graddiv2_vn: np.array, - vn: np.array, -) -> np.array: - nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) - vn = vn + (scal_divdamp + bdy_divdamp * nudgecoeff_e) * z_graddiv2_vn - return vn - - -def test_mo_solve_nonhydro_stencil_27(): - mesh = SimpleMesh() - - scal_divdamp = random_field(mesh, KDim) - bdy_divdamp = random_field(mesh, KDim) - nudgecoeff_e = random_field(mesh, EdgeDim) - z_graddiv2_vn = random_field(mesh, EdgeDim, KDim) - vn = random_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_27_numpy( - np.asarray(scal_divdamp), - np.asarray(bdy_divdamp), - np.asarray(nudgecoeff_e), - np.asarray(z_graddiv2_vn), - np.asarray(vn), - ) - mo_solve_nonhydro_stencil_27( - scal_divdamp, - bdy_divdamp, - nudgecoeff_e, - z_graddiv2_vn, - vn, - offset_provider={}, - ) - assert np.allclose(vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_28.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_28.py deleted file mode 100644 index fc43e0330e..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_28.py +++ /dev/null @@ -1,48 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_28 import ( - mo_solve_nonhydro_stencil_28, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_28_numpy( - vn_incr: np.array, vn: np.array, iau_wgt_dyn -) -> np.array: - vn = vn + (iau_wgt_dyn * vn_incr) - return vn - - -def test_mo_solve_nonhydro_stencil_28(): - mesh = SimpleMesh() - - vn_incr = random_field(mesh, EdgeDim, KDim) - vn = random_field(mesh, EdgeDim, KDim) - iau_wgt_dyn = 5.0 - - ref = mo_solve_nonhydro_stencil_28_numpy( - np.asarray(vn_incr), np.asarray(vn), iau_wgt_dyn - ) - mo_solve_nonhydro_stencil_28( - vn_incr, - vn, - iau_wgt_dyn, - offset_provider={}, - ) - assert np.allclose(vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_29.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_29.py deleted file mode 100644 index a43fc0fd10..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_29.py +++ /dev/null @@ -1,50 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_29 import ( - mo_solve_nonhydro_stencil_29, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_29_numpy( - grf_tend_vn: np.array, vn_now: np.array, dtime -) -> np.array: - vn_new = vn_now + dtime * grf_tend_vn - return vn_new - - -def test_mo_solve_nonhydro_stencil_29(): - mesh = SimpleMesh() - - grf_tend_vn = random_field(mesh, EdgeDim, KDim) - vn_now = random_field(mesh, EdgeDim, KDim) - vn_new = zero_field(mesh, EdgeDim, KDim) - dtime = 6.0 - - ref = mo_solve_nonhydro_stencil_29_numpy( - np.asarray(grf_tend_vn), np.asarray(vn_now), dtime - ) - mo_solve_nonhydro_stencil_29( - grf_tend_vn, - vn_now, - vn_new, - dtime, - offset_provider={}, - ) - assert np.allclose(vn_new, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py deleted file mode 100644 index 65a3f1ce4f..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py +++ /dev/null @@ -1,76 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_30 import ( - mo_solve_nonhydro_stencil_30, -) -from icon4py.common.dimension import E2C2EDim, E2C2EODim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_30_numpy( - e2c2e: np.array, - e2c2eO: np.array, - e_flx_avg: np.array, - vn: np.array, - geofac_grdiv: np.array, - rbf_vec_coeff_e: np.array, -) -> tuple[np.array]: - e_flx_avg = np.expand_dims(e_flx_avg, axis=-1) - z_vn_avg = np.sum(vn[e2c2eO] * e_flx_avg, axis=1) - geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) - z_graddiv_vn = np.sum(vn[e2c2eO] * geofac_grdiv, axis=1) - rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) - vt = np.sum(vn[e2c2e] * rbf_vec_coeff_e, axis=1) - return z_vn_avg, z_graddiv_vn, vt - - -def test_mo_solve_nonhydro_stencil_30_vt(): - mesh = SimpleMesh() - - e_flx_avg = random_field(mesh, EdgeDim, E2C2EODim) - geofac_grdiv = random_field(mesh, EdgeDim, E2C2EODim) - rbf_vec_coeff_e = random_field(mesh, EdgeDim, E2C2EDim) - vn = random_field(mesh, EdgeDim, KDim) - z_vn_avg = zero_field(mesh, EdgeDim, KDim) - z_graddiv_vn = zero_field(mesh, EdgeDim, KDim) - vt = zero_field(mesh, EdgeDim, KDim) - - z_vn_avg_ref, z_graddiv_vn_ref, vt_ref = mo_solve_nonhydro_stencil_30_numpy( - mesh.e2c2e, - mesh.e2c2eO, - np.asarray(e_flx_avg), - np.asarray(vn), - np.asarray(geofac_grdiv), - np.asarray(rbf_vec_coeff_e), - ) - mo_solve_nonhydro_stencil_30( - e_flx_avg, - vn, - geofac_grdiv, - rbf_vec_coeff_e, - z_vn_avg, - z_graddiv_vn, - vt, - offset_provider={ - "E2C2EO": mesh.get_e2c2eO_offset_provider(), - "E2C2E": mesh.get_e2c2e_offset_provider(), - }, - ) - assert np.allclose(z_vn_avg_ref, z_vn_avg) - assert np.allclose(z_graddiv_vn_ref, z_graddiv_vn) - assert np.allclose(vt_ref, vt) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py deleted file mode 100644 index 556edef124..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_31 import ( - mo_solve_nonhydro_stencil_31, -) -from icon4py.common.dimension import E2C2EODim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_31_numpy( - e2c2eO: np.array, e_flx_avg: np.array, vn: np.array -) -> np.array: - geofac_grdiv = np.expand_dims(e_flx_avg, axis=-1) - z_temp = np.sum(vn[e2c2eO] * geofac_grdiv, axis=1) - return z_temp - - -def test_mo_solve_nonhydro_stencil_31(): - mesh = SimpleMesh() - - e_flx_avg = random_field(mesh, EdgeDim, E2C2EODim) - vn = random_field(mesh, EdgeDim, KDim) - z_vn_avg = zero_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_31_numpy( - mesh.e2c2eO, np.asarray(e_flx_avg), np.asarray(vn) - ) - mo_solve_nonhydro_stencil_31( - e_flx_avg, - vn, - z_vn_avg, - offset_provider={"E2C2EO": mesh.get_e2c2eO_offset_provider()}, - ) - assert np.allclose(z_vn_avg, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_32.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_32.py deleted file mode 100644 index 7c716e78f2..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_32.py +++ /dev/null @@ -1,64 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_32 import ( - mo_solve_nonhydro_stencil_32, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_32_numpy( - z_rho_e: np.array, - z_vn_avg: np.array, - ddqz_z_full_e: np.array, - z_theta_v_e: np.array, -) -> tuple[np.array]: - mass_fl_e = z_rho_e * z_vn_avg * ddqz_z_full_e - z_theta_v_fl_e = mass_fl_e * z_theta_v_e - return mass_fl_e, z_theta_v_fl_e - - -def test_mo_solve_nonhydro_stencil_32(): - mesh = SimpleMesh() - - z_rho_e = random_field(mesh, EdgeDim, KDim) - z_vn_avg = random_field(mesh, EdgeDim, KDim) - ddqz_z_full_e = random_field(mesh, EdgeDim, KDim) - mass_fl_e = random_field(mesh, EdgeDim, KDim) - z_theta_v_e = random_field(mesh, EdgeDim, KDim) - z_theta_v_fl_e = zero_field(mesh, EdgeDim, KDim) - - mass_fl_e_ref, z_theta_v_fl_e_ref = mo_solve_nonhydro_stencil_32_numpy( - np.asarray(z_rho_e), - np.asarray(z_vn_avg), - np.asarray(ddqz_z_full_e), - np.asarray(z_theta_v_e), - ) - - mo_solve_nonhydro_stencil_32( - z_rho_e, - z_vn_avg, - ddqz_z_full_e, - z_theta_v_e, - mass_fl_e, - z_theta_v_fl_e, - offset_provider={}, - ) - - assert np.allclose(mass_fl_e, mass_fl_e_ref) - assert np.allclose(z_theta_v_fl_e, z_theta_v_fl_e_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_33.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_33.py deleted file mode 100644 index a9f7c3f87e..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_33.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_33 import ( - mo_solve_nonhydro_stencil_33, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_33_numpy( - vn_traj: np.array, mass_flx_me: np.array -) -> tuple[np.array]: - vn_traj = np.zeros_like(vn_traj) - mass_flx_me = np.zeros_like(mass_flx_me) - return vn_traj, mass_flx_me - - -def test_mo_solve_nonhydro_stencil_33(): - mesh = SimpleMesh() - - vn_traj = zero_field(mesh, EdgeDim, KDim) - mass_flx_me = zero_field(mesh, EdgeDim, KDim) - - vn_traj_ref, mass_flx_me_ref = mo_solve_nonhydro_stencil_33_numpy( - np.asarray(vn_traj), np.asarray(mass_flx_me) - ) - mo_solve_nonhydro_stencil_33( - vn_traj, - mass_flx_me, - offset_provider={}, - ) - - assert np.allclose(vn_traj, vn_traj_ref) - assert np.allclose(mass_flx_me, mass_flx_me_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_34.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_34.py deleted file mode 100644 index db93e47bda..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_34.py +++ /dev/null @@ -1,63 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_34 import ( - mo_solve_nonhydro_stencil_34, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_34_numpy( - z_vn_avg: np.array, - mass_fl_e: np.array, - vn_traj: np.array, - mass_flx_me: np.array, - r_nsubsteps, -) -> tuple[np.array]: - vn_traj = vn_traj + r_nsubsteps * z_vn_avg - mass_flx_me = mass_flx_me + r_nsubsteps * mass_fl_e - return vn_traj, mass_flx_me - - -def test_mo_solve_nonhydro_stencil_34(): - mesh = SimpleMesh() - - mass_fl_e = random_field(mesh, EdgeDim, KDim) - mass_flx_me = random_field(mesh, EdgeDim, KDim) - z_vn_avg = random_field(mesh, EdgeDim, KDim) - vn_traj = random_field(mesh, EdgeDim, KDim) - r_nsubsteps = 9.0 - - vn_traj_ref, mass_flx_me_ref = mo_solve_nonhydro_stencil_34_numpy( - np.asarray(z_vn_avg), - np.asarray(mass_fl_e), - np.asarray(vn_traj), - np.asarray(mass_flx_me), - r_nsubsteps, - ) - - mo_solve_nonhydro_stencil_34( - z_vn_avg, - mass_fl_e, - vn_traj, - mass_flx_me, - r_nsubsteps, - offset_provider={}, - ) - assert np.allclose(vn_traj_ref, vn_traj) - assert np.allclose(mass_flx_me_ref, mass_flx_me) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_35.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_35.py deleted file mode 100644 index a8403a7dac..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_35.py +++ /dev/null @@ -1,52 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_35 import ( - mo_solve_nonhydro_stencil_35, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_35_numpy( - vn: np.array, ddxn_z_full: np.array, ddxt_z_full: np.array, vt: np.array -) -> np.array: - z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full - return z_w_concorr_me - - -def test_mo_solve_nonhydro_stencil_35(): - mesh = SimpleMesh() - - vn = random_field(mesh, EdgeDim, KDim) - ddxn_z_full = random_field(mesh, EdgeDim, KDim) - ddxt_z_full = random_field(mesh, EdgeDim, KDim) - vt = random_field(mesh, EdgeDim, KDim) - z_w_concorr_me = zero_field(mesh, EdgeDim, KDim) - - ref = mo_solve_nonhydro_stencil_35_numpy( - np.asarray(vn), np.asarray(ddxn_z_full), np.asarray(ddxt_z_full), np.asarray(vt) - ) - mo_solve_nonhydro_stencil_35( - vn, - ddxn_z_full, - ddxt_z_full, - vt, - z_w_concorr_me, - offset_provider={}, - ) - assert np.allclose(z_w_concorr_me, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_37.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_37.py deleted file mode 100644 index 6ad99a0413..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_37.py +++ /dev/null @@ -1,54 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_37 import ( - mo_solve_nonhydro_stencil_37, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_37_numpy(vn: np.array, vt: np.array) -> tuple[np.array]: - vn_ie = vn - z_vt_ie = vt - z_kin_hor_e = 0.5 * (pow(vn, 2) + pow(vt, 2)) - return vn_ie, z_vt_ie, z_kin_hor_e - - -def test_mo_solve_nonhydro_stencil_37_z_kin_hor_e(): - mesh = SimpleMesh() - - vt = random_field(mesh, EdgeDim, KDim) - vn = random_field(mesh, EdgeDim, KDim) - vn_ie = zero_field(mesh, EdgeDim, KDim) - z_kin_hor_e = zero_field(mesh, EdgeDim, KDim) - z_vt_ie = zero_field(mesh, EdgeDim, KDim) - - vn_ie_ref, z_vt_ie_ref, z_kin_hor_e_ref = mo_solve_nonhydro_stencil_37_numpy( - np.asarray(vn), np.asarray(vt) - ) - mo_solve_nonhydro_stencil_37( - vn, - vt, - vn_ie, - z_vt_ie, - z_kin_hor_e, - offset_provider={}, - ) - assert np.allclose(z_kin_hor_e_ref, z_kin_hor_e) - assert np.allclose(vn_ie_ref, vn_ie) - assert np.allclose(z_vt_ie_ref, z_vt_ie) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py deleted file mode 100644 index 2dbffe70dd..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py +++ /dev/null @@ -1,59 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_41 import ( - mo_solve_nonhydro_stencil_41, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_41_numpy( - c2e: np.array, geofac_div: np.array, mass_fl_e: np.array, z_theta_v_fl_e: np.array -) -> tuple[np.array]: - geofac_div = np.expand_dims(geofac_div, axis=-1) - z_flxdiv_mass = np.sum(geofac_div * mass_fl_e[c2e], axis=1) - z_flxdiv_theta = np.sum(geofac_div * z_theta_v_fl_e[c2e], axis=1) - return z_flxdiv_mass, z_flxdiv_theta - - -def test_mo_solve_nonhydro_stencil_41_z_flxdiv_theta(): - mesh = SimpleMesh() - - geofac_div = random_field(mesh, CellDim, C2EDim) - z_theta_v_fl_e = random_field(mesh, EdgeDim, KDim) - z_flxdiv_theta = zero_field(mesh, CellDim, KDim) - mass_fl_e = random_field(mesh, EdgeDim, KDim) - z_flxdiv_mass = zero_field(mesh, CellDim, KDim) - - z_flxdiv_mass_ref, z_flxdiv_theta_ref = mo_solve_nonhydro_stencil_41_numpy( - mesh.c2e, - np.asarray(geofac_div), - np.asarray(mass_fl_e), - np.asarray(z_theta_v_fl_e), - ) - mo_solve_nonhydro_stencil_41( - geofac_div, - mass_fl_e, - z_theta_v_fl_e, - z_flxdiv_mass, - z_flxdiv_theta, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, - ) - - assert np.allclose(z_flxdiv_mass, z_flxdiv_mass_ref) - assert np.allclose(z_flxdiv_theta, z_flxdiv_theta_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_42.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_42.py deleted file mode 100644 index c27983112d..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_42.py +++ /dev/null @@ -1,96 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_42 import ( - mo_solve_nonhydro_stencil_42, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_42_numpy( - w_nnow: np.array, - ddt_w_adv_ntl1: np.array, - ddt_w_adv_ntl2: np.array, - z_th_ddz_exner_c: np.array, - rho_ic: np.array, - w_concorr_c: np.array, - vwind_expl_wgt: np.array, - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, -) -> tuple[np.array]: - z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c - ) - vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) - z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) - return z_w_expl, z_contr_w_fl_l - - -def test_mo_solve_nonhydro_stencil_42(): - mesh = SimpleMesh() - - w_nnow = random_field(mesh, CellDim, KDim) - ddt_w_adv_ntl1 = random_field(mesh, CellDim, KDim) - ddt_w_adv_ntl2 = random_field(mesh, CellDim, KDim) - z_th_ddz_exner_c = random_field(mesh, CellDim, KDim) - z_w_expl = zero_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim) - w_concorr_c = random_field(mesh, CellDim, KDim) - vwind_expl_wgt = random_field(mesh, CellDim) - z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) - dtime = 5.0 - wgt_nnow_vel = 8.0 - wgt_nnew_vel = 9.0 - cpd = 10.0 - - z_w_expl_ref, z_contr_w_fl_l_ref = mo_solve_nonhydro_stencil_42_numpy( - np.asarray(w_nnow), - np.asarray(ddt_w_adv_ntl1), - np.asarray(ddt_w_adv_ntl2), - np.asarray(z_th_ddz_exner_c), - np.asarray(rho_ic), - np.asarray(w_concorr_c), - np.asarray(vwind_expl_wgt), - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, - ) - mo_solve_nonhydro_stencil_42( - z_w_expl, - w_nnow, - ddt_w_adv_ntl1, - ddt_w_adv_ntl2, - z_th_ddz_exner_c, - z_contr_w_fl_l, - rho_ic, - w_concorr_c, - vwind_expl_wgt, - dtime, - wgt_nnow_vel, - wgt_nnew_vel, - cpd, - offset_provider={}, - ) - - assert np.allclose(z_w_expl, z_w_expl_ref) - assert np.allclose(z_contr_w_fl_l, z_contr_w_fl_l_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_43.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_43.py deleted file mode 100644 index 238997182b..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_43.py +++ /dev/null @@ -1,80 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_43 import ( - mo_solve_nonhydro_stencil_43, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_43_numpy( - w_nnow: np.array, - ddt_w_adv_ntl1: np.array, - z_th_ddz_exner_c: np.array, - rho_ic: np.array, - w_concorr_c: np.array, - vwind_expl_wgt: np.array, - dtime: float, - cpd: float, -) -> tuple[np.array]: - vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, -1) - z_w_expl = w_nnow + dtime * (ddt_w_adv_ntl1 - cpd * z_th_ddz_exner_c) - z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) - return z_w_expl, z_contr_w_fl_l - - -def test_mo_solve_nonhydro_stencil_43(): - mesh = SimpleMesh() - - w_nnow = random_field(mesh, CellDim, KDim) - ddt_w_adv_ntl1 = random_field(mesh, CellDim, KDim) - z_th_ddz_exner_c = random_field(mesh, CellDim, KDim) - z_w_expl = zero_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim) - w_concorr_c = random_field(mesh, CellDim, KDim) - vwind_expl_wgt = random_field(mesh, CellDim) - z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) - dtime = 5.0 - cpd = 10.0 - - z_w_expl_ref, z_contr_w_fl_l_ref = mo_solve_nonhydro_stencil_43_numpy( - np.asarray(w_nnow), - np.asarray(ddt_w_adv_ntl1), - np.asarray(z_th_ddz_exner_c), - np.asarray(rho_ic), - np.asarray(w_concorr_c), - np.asarray(vwind_expl_wgt), - dtime, - cpd, - ) - mo_solve_nonhydro_stencil_43( - z_w_expl, - w_nnow, - ddt_w_adv_ntl1, - z_th_ddz_exner_c, - z_contr_w_fl_l, - rho_ic, - w_concorr_c, - vwind_expl_wgt, - dtime, - cpd, - offset_provider={}, - ) - - assert np.allclose(z_w_expl, z_w_expl_ref) - assert np.allclose(z_contr_w_fl_l, z_contr_w_fl_l_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_44.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_44.py deleted file mode 100644 index 721b811fb1..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_44.py +++ /dev/null @@ -1,89 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_44 import ( - mo_solve_nonhydro_stencil_44, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_44_numpy( - exner_nnow: np.array, - rho_nnow: np.array, - theta_v_nnow: np.array, - inv_ddqz_z_full: np.array, - vwind_impl_wgt: np.array, - theta_v_ic: np.array, - rho_ic: np.array, - dtime, - rd, - cvd, -) -> tuple[np.array]: - z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full - - vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic - return z_beta, z_alpha - - -def test_mo_solve_nonhydro_stencil_44_z_alpha(): - mesh = SimpleMesh() - - exner_nnow = random_field(mesh, CellDim, KDim) - rho_nnow = random_field(mesh, CellDim, KDim) - theta_v_nnow = random_field(mesh, CellDim, KDim) - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - vwind_impl_wgt = random_field(mesh, CellDim) - theta_v_ic = random_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim) - z_alpha = zero_field(mesh, CellDim, KDim) - z_beta = zero_field(mesh, CellDim, KDim) - dtime = 10.0 - rd = 5.0 - cvd = 3.0 - - z_beta_ref, z_alpha_ref = mo_solve_nonhydro_stencil_44_numpy( - np.asarray(exner_nnow), - np.asarray(rho_nnow), - np.asarray(theta_v_nnow), - np.asarray(inv_ddqz_z_full), - np.asarray(vwind_impl_wgt), - np.asarray(theta_v_ic), - np.asarray(rho_ic), - dtime, - rd, - cvd, - ) - - mo_solve_nonhydro_stencil_44( - z_beta, - exner_nnow, - rho_nnow, - theta_v_nnow, - inv_ddqz_z_full, - z_alpha, - vwind_impl_wgt, - theta_v_ic, - rho_ic, - dtime, - rd, - cvd, - offset_provider={}, - ) - assert np.allclose(z_beta_ref, z_beta) - assert np.allclose(z_alpha_ref, z_alpha) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_45.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_45.py deleted file mode 100644 index 7deb2c25de..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_45.py +++ /dev/null @@ -1,40 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_45 import ( - mo_solve_nonhydro_stencil_45, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_45_numpy(z_alpha: np.array) -> np.array: - z_alpha = np.zeros_like(z_alpha) - return z_alpha - - -def test_mo_solve_nonhydro_stencil_45(): - mesh = SimpleMesh() - - z_alpha = zero_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_45_numpy(np.asarray(z_alpha)) - mo_solve_nonhydro_stencil_45( - z_alpha, - offset_provider={}, - ) - assert np.allclose(z_alpha, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_46.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_46.py deleted file mode 100644 index 3cd93384fd..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_46.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_46 import ( - mo_solve_nonhydro_stencil_46, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_46_numpy( - w_nnew: np.array, - z_contr_w_fl_l: np.array, -) -> tuple[np.array]: - w_nnew = np.zeros_like(w_nnew) - z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) - return w_nnew, z_contr_w_fl_l - - -def test_mo_solve_nonhydro_stencil_46_z_contr_w_fl_l(): - mesh = SimpleMesh() - - z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) - w_nnew = zero_field(mesh, CellDim, KDim) - - w_nnew_ref, z_contr_w_fl_l_ref = mo_solve_nonhydro_stencil_46_numpy( - np.asarray(w_nnew), np.asarray(z_contr_w_fl_l) - ) - mo_solve_nonhydro_stencil_46( - w_nnew, - z_contr_w_fl_l, - offset_provider={}, - ) - assert np.allclose(w_nnew_ref, w_nnew) - assert np.allclose(z_contr_w_fl_l_ref, z_contr_w_fl_l) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_47.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_47.py deleted file mode 100644 index bbb073f7e3..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_47.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_47 import ( - mo_solve_nonhydro_stencil_47, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_47_numpy( - w_concorr_c: np.array, - z_contr_w_fl_l: np.array, -) -> tuple[np.array]: - w_nnew = w_concorr_c - z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) - return w_nnew, z_contr_w_fl_l - - -def test_mo_solve_nonhydro_stencil_47_z_contr_w_fl_l(): - mesh = SimpleMesh() - - w_concorr_c = random_field(mesh, CellDim, KDim) - z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) - w_nnew = zero_field(mesh, CellDim, KDim) - - w_nnew_ref, z_contr_w_fl_l_ref = mo_solve_nonhydro_stencil_47_numpy( - np.asarray(w_concorr_c), np.asarray(z_contr_w_fl_l) - ) - mo_solve_nonhydro_stencil_47( - w_nnew, - z_contr_w_fl_l, - w_concorr_c, - offset_provider={}, - ) - assert np.allclose(w_nnew, w_nnew_ref) - assert np.allclose(z_contr_w_fl_l, z_contr_w_fl_l_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_48.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_48.py deleted file mode 100644 index 363a93bd1f..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_48.py +++ /dev/null @@ -1,101 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_48 import ( - mo_solve_nonhydro_stencil_48, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_48_numpy( - rho_nnow: np.array, - inv_ddqz_z_full: np.array, - z_flxdiv_mass: np.array, - z_contr_w_fl_l: np.array, - exner_pr: np.array, - z_beta: np.array, - z_flxdiv_theta: np.array, - theta_v_ic: np.array, - ddt_exner_phy: np.array, - dtime, -) -> tuple[np.array, np.array]: - z_rho_expl = rho_nnow - dtime * inv_ddqz_z_full * ( - z_flxdiv_mass + z_contr_w_fl_l[:, :-1] - z_contr_w_fl_l[:, 1:] - ) - - z_exner_expl = ( - exner_pr - - z_beta - * ( - z_flxdiv_theta - + (theta_v_ic * z_contr_w_fl_l)[:, :-1] - - (theta_v_ic * z_contr_w_fl_l)[:, 1:] - ) - + dtime * ddt_exner_phy - ) - return z_rho_expl, z_exner_expl - - -def test_mo_solve_nonhydro_stencil_48(): - mesh = SimpleMesh() - - dtime = 1.0 - rho_nnow = random_field(mesh, CellDim, KDim) - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - z_flxdiv_mass = random_field(mesh, CellDim, KDim) - z_contr_w_fl_l = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - exner_pr = random_field(mesh, CellDim, KDim) - z_beta = random_field(mesh, CellDim, KDim) - z_flxdiv_theta = random_field(mesh, CellDim, KDim) - theta_v_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - ddt_exner_phy = random_field(mesh, CellDim, KDim) - - z_rho_expl = zero_field(mesh, CellDim, KDim) - z_exner_expl = zero_field(mesh, CellDim, KDim) - - z_rho_expl_ref, z_exner_expl_ref = mo_solve_nonhydro_stencil_48_numpy( - np.asarray(rho_nnow), - np.asarray(inv_ddqz_z_full), - np.asarray(z_flxdiv_mass), - np.asarray(z_contr_w_fl_l), - np.asarray(exner_pr), - np.asarray(z_beta), - np.asarray(z_flxdiv_theta), - np.asarray(theta_v_ic), - np.asarray(ddt_exner_phy), - dtime, - ) - - mo_solve_nonhydro_stencil_48( - z_rho_expl, - z_exner_expl, - rho_nnow, - inv_ddqz_z_full, - z_flxdiv_mass, - z_contr_w_fl_l, - exner_pr, - z_beta, - z_flxdiv_theta, - theta_v_ic, - ddt_exner_phy, - dtime, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_rho_expl_ref, z_rho_expl) - assert np.allclose(z_exner_expl_ref, z_exner_expl) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_49.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_49.py deleted file mode 100644 index 6b6c84c0d3..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_49.py +++ /dev/null @@ -1,101 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_49 import ( - mo_solve_nonhydro_stencil_49, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_49_numpy( - rho_nnow: np.array, - inv_ddqz_z_full: np.array, - z_flxdiv_mass: np.array, - z_contr_w_fl_l: np.array, - exner_pr: np.array, - z_beta: np.array, - z_flxdiv_theta: np.array, - theta_v_ic: np.array, - ddt_exner_phy: np.array, - dtime, -) -> tuple[np.array]: - z_rho_expl = rho_nnow - dtime * inv_ddqz_z_full * ( - z_flxdiv_mass + z_contr_w_fl_l[:, :-1] - z_contr_w_fl_l[:, 1:] - ) - z_exner_expl = ( - exner_pr - - z_beta - * ( - z_flxdiv_theta - + (theta_v_ic * z_contr_w_fl_l)[:, :-1] - - (theta_v_ic * z_contr_w_fl_l)[:, 1:] - ) - + dtime * ddt_exner_phy - ) - return z_rho_expl, z_exner_expl - - -def test_mo_solve_nonhydro_stencil_49(): - mesh = SimpleMesh() - - dtime = 7.0 - - rho_nnow = random_field(mesh, CellDim, KDim) - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - z_flxdiv_mass = random_field(mesh, CellDim, KDim) - z_contr_w_fl_l = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - exner_pr = random_field(mesh, CellDim, KDim) - z_beta = random_field(mesh, CellDim, KDim) - z_flxdiv_theta = random_field(mesh, CellDim, KDim) - theta_v_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - ddt_exner_phy = random_field(mesh, CellDim, KDim) - - z_rho_expl = zero_field(mesh, CellDim, KDim) - z_exner_expl = zero_field(mesh, CellDim, KDim) - - z_rho_expl_ref, z_exner_expl_ref = mo_solve_nonhydro_stencil_49_numpy( - np.asarray(rho_nnow), - np.asarray(inv_ddqz_z_full), - np.asarray(z_flxdiv_mass), - np.asarray(z_contr_w_fl_l), - np.asarray(exner_pr), - np.asarray(z_beta), - np.asarray(z_flxdiv_theta), - np.asarray(theta_v_ic), - np.asarray(ddt_exner_phy), - dtime, - ) - - mo_solve_nonhydro_stencil_49( - z_rho_expl, - z_exner_expl, - rho_nnow, - inv_ddqz_z_full, - z_flxdiv_mass, - z_contr_w_fl_l, - exner_pr, - z_beta, - z_flxdiv_theta, - theta_v_ic, - ddt_exner_phy, - dtime, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_rho_expl[:, :-1], z_rho_expl_ref[:, :-1]) - assert np.allclose(z_exner_expl[:, :-1], z_exner_expl_ref[:, :-1]) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_50.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_50.py deleted file mode 100644 index bbe4e81929..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_50.py +++ /dev/null @@ -1,64 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_50 import ( - mo_solve_nonhydro_stencil_50, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_50_numpy( - z_rho_expl: np.array, - rho_incr: np.array, - z_exner_expl: np.array, - exner_incr: np.array, - iau_wgt_dyn, -) -> tuple[np.array]: - z_rho_expl = z_rho_expl + iau_wgt_dyn * rho_incr - z_exner_expl = z_exner_expl + iau_wgt_dyn * exner_incr - return z_rho_expl, z_exner_expl - - -def test_mo_solve_nonhydro_stencil_50(): - mesh = SimpleMesh() - - z_exner_expl = random_field(mesh, CellDim, KDim) - exner_incr = random_field(mesh, CellDim, KDim) - z_rho_expl = random_field(mesh, CellDim, KDim) - rho_incr = random_field(mesh, CellDim, KDim) - iau_wgt_dyn = 8.0 - - z_rho_expl_ref, z_exner_expl_ref = mo_solve_nonhydro_stencil_50_numpy( - np.asarray(z_rho_expl), - np.asarray(rho_incr), - np.asarray(z_exner_expl), - np.asarray(exner_incr), - iau_wgt_dyn, - ) - - mo_solve_nonhydro_stencil_50( - z_rho_expl, - z_exner_expl, - rho_incr, - exner_incr, - iau_wgt_dyn, - offset_provider={}, - ) - - assert np.allclose(z_rho_expl, z_rho_expl_ref) - assert np.allclose(z_exner_expl, z_exner_expl_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_53.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_53.py deleted file mode 100644 index 100eb1d575..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_53.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_53 import ( - mo_solve_nonhydro_stencil_53, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_53_numpy( - z_q: np.array, - w: np.array, -) -> np.array: - w_new = np.zeros_like(w) - last_k_level = w.shape[1] - 1 - - w_new[:, last_k_level] = w[:, last_k_level] - for k in reversed(range(1, last_k_level)): - w_new[:, k] = w[:, k] + w_new[:, k + 1] * z_q[:, k] - w_new[:, 0] = w[:, 0] - return w_new - - -def test_mo_solve_nonhydro_stencil_53(): - mesh = SimpleMesh() - z_q = random_field(mesh, CellDim, KDim) - w = random_field(mesh, CellDim, KDim) - - w_ref = mo_solve_nonhydro_stencil_53_numpy( - np.asarray(z_q), - np.asarray(w), - ) - - mo_solve_nonhydro_stencil_53(z_q, w, offset_provider={"Koff": KDim}) - - assert np.allclose(w_ref, w) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_54.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_54.py deleted file mode 100644 index 6eead35fed..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_54.py +++ /dev/null @@ -1,50 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_54 import ( - mo_solve_nonhydro_stencil_54, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_54_numpy( - z_raylfac: np.array, w_1: np.array, w: np.array -) -> np.array: - z_raylfac = np.expand_dims(z_raylfac, axis=0) - w_1 = np.expand_dims(w_1, axis=-1) - w = z_raylfac * w + (1.0 - z_raylfac) * w_1 - return w - - -def test_mo_solve_nonhydro_stencil_54(): - mesh = SimpleMesh() - - z_raylfac = random_field(mesh, KDim) - w_1 = random_field(mesh, CellDim) - w = random_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_54_numpy( - np.asarray(z_raylfac), np.asarray(w_1), np.asarray(w) - ) - mo_solve_nonhydro_stencil_54( - z_raylfac, - w_1, - w, - offset_provider={}, - ) - assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_55.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_55.py deleted file mode 100644 index 2d561f9b9e..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_55.py +++ /dev/null @@ -1,124 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_55 import ( - mo_solve_nonhydro_stencil_55, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_55_numpy( - z_rho_expl: np.array, - vwind_impl_wgt: np.array, - inv_ddqz_z_full: np.array, - rho_ic: np.array, - w: np.array, - z_exner_expl: np.array, - exner_ref_mc: np.array, - z_alpha: np.array, - z_beta: np.array, - rho_now: np.array, - theta_v_now: np.array, - exner_now: np.array, - dtime, - cvd_o_rd, -) -> tuple[np.array]: - rho_ic_offset_1 = rho_ic[:, 1:] - w_offset_0 = w[:, :-1] - w_offset_1 = w[:, 1:] - z_alpha_offset_1 = z_alpha[:, 1:] - vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=1) - rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( - rho_ic[:, :-1] * w_offset_0 - rho_ic_offset_1 * w_offset_1 - ) - exner_new = ( - z_exner_expl - + exner_ref_mc - - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) - ) - theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new - ) - return rho_new, exner_new, theta_v_new - - -def test_mo_solve_nonhydro_stencil_55(): - mesh = SimpleMesh() - - z_rho_expl = random_field(mesh, CellDim, KDim) - vwind_impl_wgt = random_field(mesh, CellDim) - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - w = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - z_exner_expl = random_field(mesh, CellDim, KDim) - exner_ref_mc = random_field(mesh, CellDim, KDim) - z_alpha = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - z_beta = random_field(mesh, CellDim, KDim) - rho_now = random_field(mesh, CellDim, KDim) - theta_v_now = random_field(mesh, CellDim, KDim) - exner_now = random_field(mesh, CellDim, KDim) - rho_new = zero_field(mesh, CellDim, KDim) - exner_new = zero_field(mesh, CellDim, KDim) - theta_v_new = zero_field(mesh, CellDim, KDim) - dtime = 5.0 - cvd_o_rd = 9.0 - - rho_new_ref, exner_new_ref, theta_v_new_ref = mo_solve_nonhydro_stencil_55_numpy( - np.asarray(z_rho_expl), - np.asarray(vwind_impl_wgt), - np.asarray(inv_ddqz_z_full), - np.asarray(rho_ic), - np.asarray(w), - np.asarray(z_exner_expl), - np.asarray(exner_ref_mc), - np.asarray(z_alpha), - np.asarray(z_beta), - np.asarray(rho_now), - np.asarray(theta_v_now), - np.asarray(exner_now), - dtime, - cvd_o_rd, - ) - - mo_solve_nonhydro_stencil_55( - z_rho_expl, - vwind_impl_wgt, - inv_ddqz_z_full, - rho_ic, - w, - z_exner_expl, - exner_ref_mc, - z_alpha, - z_beta, - rho_now, - theta_v_now, - exner_now, - rho_new, - exner_new, - theta_v_new, - dtime, - cvd_o_rd, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(rho_new, rho_new_ref) - assert np.allclose(exner_new, exner_new_ref) - assert np.allclose(theta_v_new, theta_v_new_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_56_63.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_56_63.py deleted file mode 100644 index eb9521b275..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_56_63.py +++ /dev/null @@ -1,53 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_56_63 import ( - mo_solve_nonhydro_stencil_56_63, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_56_63_numpy( - inv_ddqz_z_full: np.array, - w: np.array, - w_concorr_c: np.array, -) -> np.array: - z_dwdz_dd = inv_ddqz_z_full * ( - (w[:, :-1] - w[:, 1:]) - (w_concorr_c[:, :-1] - w_concorr_c[:, 1:]) - ) - return z_dwdz_dd - - -def test_mo_solve_nonhydro_stencil_56_63(): - mesh = SimpleMesh() - inv_ddqz_z_full = random_field(mesh, CellDim, KDim) - w = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - w_concorr_c = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - z_dwdz_dd = random_field(mesh, CellDim, KDim) - - z_dwdz_dd_ref = mo_solve_nonhydro_stencil_56_63_numpy( - np.asarray(inv_ddqz_z_full), - np.asarray(w), - np.asarray(w_concorr_c), - ) - - mo_solve_nonhydro_stencil_56_63( - inv_ddqz_z_full, w, w_concorr_c, z_dwdz_dd, offset_provider={"Koff": KDim} - ) - - assert np.allclose(z_dwdz_dd_ref, z_dwdz_dd) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_57.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_57.py deleted file mode 100644 index 4f829406a9..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_57.py +++ /dev/null @@ -1,40 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_57 import ( - mo_solve_nonhydro_stencil_57, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_57_numpy(mass_flx_ic: np.array) -> np.array: - mass_flx_ic = np.zeros_like(mass_flx_ic) - return mass_flx_ic - - -def test_mo_solve_nonhydro_stencil_57(): - mesh = SimpleMesh() - - mass_flx_ic = zero_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_57_numpy(np.asarray(mass_flx_ic)) - mo_solve_nonhydro_stencil_57( - mass_flx_ic, - offset_provider={}, - ) - assert np.allclose(mass_flx_ic, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_58.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_58.py deleted file mode 100644 index 5065b54823..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_58.py +++ /dev/null @@ -1,68 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_58 import ( - mo_solve_nonhydro_stencil_58, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_58_numpy( - z_contr_w_fl_l: np.array, - rho_ic: np.array, - vwind_impl_wgt: np.array, - w: np.array, - mass_flx_ic: np.array, - r_nsubsteps, -) -> np.array: - vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) - return mass_flx_ic - - -def test_mo_solve_nonhydro_stencil_58(): - mesh = SimpleMesh() - - z_contr_w_fl_l = random_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim) - vwind_impl_wgt = random_field(mesh, CellDim) - w = random_field(mesh, CellDim, KDim) - mass_flx_ic = random_field(mesh, CellDim, KDim) - r_nsubsteps = 7.0 - - ref = mo_solve_nonhydro_stencil_58_numpy( - np.asarray(z_contr_w_fl_l), - np.asarray(rho_ic), - np.asarray(vwind_impl_wgt), - np.asarray(w), - np.asarray(mass_flx_ic), - r_nsubsteps, - ) - - mo_solve_nonhydro_stencil_58( - z_contr_w_fl_l, - rho_ic, - vwind_impl_wgt, - w, - mass_flx_ic, - r_nsubsteps, - offset_provider={}, - ) - assert np.allclose(mass_flx_ic, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_59.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_59.py deleted file mode 100644 index e44ca54117..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_59.py +++ /dev/null @@ -1,45 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_59 import ( - mo_solve_nonhydro_stencil_59, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_59_numpy(exner: np.array) -> np.array: - exner_dyn_incr = exner - return exner_dyn_incr - - -def test_mo_solve_nonhydro_stencil_59(): - mesh = SimpleMesh() - - exner = random_field(mesh, CellDim, KDim) - exner_dyn_incr = zero_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_59_numpy( - np.asarray(exner), - ) - - mo_solve_nonhydro_stencil_59( - exner, - exner_dyn_incr, - offset_provider={}, - ) - assert np.allclose(exner_dyn_incr, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_60.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_60.py deleted file mode 100644 index dd668d408d..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_60.py +++ /dev/null @@ -1,62 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_60 import ( - mo_solve_nonhydro_stencil_60, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_60_numpy( - exner: np.array, - ddt_exner_phy: np.array, - exner_dyn_incr: np.array, - ndyn_substeps_var: float, - dtime: float, -) -> np.array: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) - return exner_dyn_incr - - -def test_mo_solve_nonhydro_stencil_60(): - mesh = SimpleMesh() - - ndyn_substeps_var, dtime = 10.0, 12.0 - exner = random_field(mesh, CellDim, KDim) - ddt_exner_phy = random_field(mesh, CellDim, KDim) - exner_dyn_incr = random_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_60_numpy( - np.asarray(exner), - np.asarray(ddt_exner_phy), - np.asarray(exner_dyn_incr), - ndyn_substeps_var, - dtime, - ) - - mo_solve_nonhydro_stencil_60( - exner, - ddt_exner_phy, - exner_dyn_incr, - ndyn_substeps_var, - dtime, - offset_provider={}, - ) - assert np.allclose(exner_dyn_incr, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_61.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_61.py deleted file mode 100644 index 1013540638..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_61.py +++ /dev/null @@ -1,78 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_61 import ( - mo_solve_nonhydro_stencil_61, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_61_numpy( - rho_now: np.array, - grf_tend_rho: np.array, - theta_v_now: np.array, - grf_tend_thv: np.array, - w_now: np.array, - grf_tend_w: np.array, - dtime, -) -> tuple[np.array]: - rho_new = rho_now + dtime * grf_tend_rho - exner_new = theta_v_now + dtime * grf_tend_thv - w_new = w_now + dtime * grf_tend_w - return rho_new, exner_new, w_new - - -def test_mo_solve_nonhydro_stencil_61(): - mesh = SimpleMesh() - - rho_now = random_field(mesh, CellDim, KDim) - grf_tend_rho = random_field(mesh, CellDim, KDim) - theta_v_now = random_field(mesh, CellDim, KDim) - grf_tend_thv = random_field(mesh, CellDim, KDim) - w_now = random_field(mesh, CellDim, KDim) - grf_tend_w = random_field(mesh, CellDim, KDim) - dtime = 5.0 - rho_new = zero_field(mesh, CellDim, KDim) - exner_new = zero_field(mesh, CellDim, KDim) - w_new = zero_field(mesh, CellDim, KDim) - - rho_new_ref, exner_new_ref, w_new_ref = mo_solve_nonhydro_stencil_61_numpy( - np.asarray(rho_now), - np.asarray(grf_tend_rho), - np.asarray(theta_v_now), - np.asarray(grf_tend_thv), - np.asarray(w_now), - np.asarray(grf_tend_w), - dtime, - ) - mo_solve_nonhydro_stencil_61( - rho_now, - grf_tend_rho, - theta_v_now, - grf_tend_thv, - w_now, - grf_tend_w, - rho_new, - exner_new, - w_new, - dtime, - offset_provider={}, - ) - assert np.allclose(rho_new, rho_new_ref) - assert np.allclose(exner_new, exner_new_ref) - assert np.allclose(w_new, w_new_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_62.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_62.py deleted file mode 100644 index 6a1d35aebc..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_62.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_62 import ( - mo_solve_nonhydro_stencil_62, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_62_numpy( - w_now: np.array, grf_tend_w: np.array, dtime: float -) -> np.array: - w_new = w_now + dtime * grf_tend_w - return w_new - - -def test_mo_solve_nonhydro_stencil_62(): - mesh = SimpleMesh() - - dtime = 10.0 - w_now = random_field(mesh, CellDim, KDim) - grf_tend_w = random_field(mesh, CellDim, KDim) - w_new = zero_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_62_numpy( - np.asarray(w_now), np.asarray(grf_tend_w), dtime - ) - - mo_solve_nonhydro_stencil_62( - w_now, - grf_tend_w, - w_new, - dtime, - offset_provider={}, - ) - assert np.allclose(w_new, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_64.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_64.py deleted file mode 100644 index e670bd1044..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_64.py +++ /dev/null @@ -1,41 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_64 import ( - mo_solve_nonhydro_stencil_64, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_64_numpy(mass_flx_ic: np.array) -> np.array: - mass_flx_ic = np.zeros_like(mass_flx_ic) - return mass_flx_ic - - -def test_mo_solve_nonhydro_stencil_64(): - mesh = SimpleMesh() - - mass_flx_ic = zero_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_64_numpy(np.asarray(mass_flx_ic)) - - mo_solve_nonhydro_stencil_64( - mass_flx_ic, - offset_provider={}, - ) - assert np.allclose(mass_flx_ic, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_65.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_65.py deleted file mode 100644 index dad4a29f94..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_65.py +++ /dev/null @@ -1,79 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_65 import ( - mo_solve_nonhydro_stencil_65, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_65_numpy( - rho_ic: np.array, - vwind_expl_wgt: np.array, - vwind_impl_wgt: np.array, - w_now: np.array, - w_new: np.array, - w_concorr_c: np.array, - mass_flx_ic: np.array, - r_nsubsteps: float, -) -> np.array: - vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) - vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) - ) - return mass_flx_ic - - -def test_mo_solve_nonhydro_stencil_65(): - mesh = SimpleMesh() - - r_nsubsteps = 10.0 - rho_ic = random_field(mesh, CellDim, KDim) - vwind_expl_wgt = random_field(mesh, CellDim) - vwind_impl_wgt = random_field(mesh, CellDim) - w_now = random_field(mesh, CellDim, KDim) - w_new = random_field(mesh, CellDim, KDim) - w_concorr_c = random_field(mesh, CellDim, KDim) - mass_flx_ic = random_field(mesh, CellDim, KDim) - - ref = mo_solve_nonhydro_stencil_65_numpy( - np.asarray(rho_ic), - np.asarray(vwind_expl_wgt), - np.asarray(vwind_impl_wgt), - np.asarray(w_now), - np.asarray(w_new), - np.asarray(w_concorr_c), - np.asarray(mass_flx_ic), - r_nsubsteps, - ) - - mo_solve_nonhydro_stencil_65( - rho_ic, - vwind_expl_wgt, - vwind_impl_wgt, - w_now, - w_new, - w_concorr_c, - mass_flx_ic, - r_nsubsteps, - offset_provider={}, - ) - assert np.allclose(mass_flx_ic, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_66.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_66.py deleted file mode 100644 index f5c35396ab..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_66.py +++ /dev/null @@ -1,73 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_66 import ( - mo_solve_nonhydro_stencil_66, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, random_mask -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_66_numpy( - bdy_halo_c: np.array, - rho: np.array, - theta_v: np.array, - exner: np.array, - rd_o_cvd: float, - rd_o_p0ref: float, -) -> tuple[np.array]: - bdy_halo_c = np.expand_dims(bdy_halo_c, axis=-1) - - theta_v = np.where(bdy_halo_c == 1, exner, theta_v) - exner = np.where( - bdy_halo_c == 1, np.exp(rd_o_cvd * np.log(rd_o_p0ref * rho * exner)), exner - ) - - return theta_v, exner - - -def test_mo_solve_nonhydro_stencil_66(): - mesh = SimpleMesh() - - rd_o_cvd = 10.0 - rd_o_p0ref = 20.0 - bdy_halo_c = random_mask(mesh, CellDim) - exner = random_field(mesh, CellDim, KDim, low=1, high=2) - rho = random_field(mesh, CellDim, KDim, low=1, high=2) - theta_v = random_field(mesh, CellDim, KDim, low=1, high=2) - - theta_v_ref, exner_ref = mo_solve_nonhydro_stencil_66_numpy( - np.asarray(bdy_halo_c), - np.asarray(rho), - np.asarray(theta_v), - np.asarray(exner), - rd_o_cvd, - rd_o_p0ref, - ) - - mo_solve_nonhydro_stencil_66( - bdy_halo_c, - rho, - theta_v, - exner, - rd_o_cvd, - rd_o_p0ref, - offset_provider={}, - ) - - assert np.allclose(theta_v, theta_v_ref) - assert np.allclose(exner, exner_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_67.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_67.py deleted file mode 100644 index 0e7cfcd884..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_67.py +++ /dev/null @@ -1,58 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_67 import ( - mo_solve_nonhydro_stencil_67, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_67_numpy( - rho: np.array, exner: np.array, rd_o_cvd: float, rd_o_p0ref: float -) -> tuple[np.array]: - - theta_v = np.copy(exner) - exner = np.exp(rd_o_cvd * np.log(rd_o_p0ref * rho * theta_v)) - - return theta_v, exner - - -def test_mo_solve_nonhydro_stencil_67(): - mesh = SimpleMesh() - - rd_o_cvd = 10.0 - rd_o_p0ref = 20.0 - rho = random_field(mesh, CellDim, KDim, low=1, high=2) - theta_v = random_field(mesh, CellDim, KDim, low=1, high=2) - exner = random_field(mesh, CellDim, KDim, low=1, high=2) - - theta_v_ref, exner_ref = mo_solve_nonhydro_stencil_67_numpy( - np.asarray(rho), np.asarray(exner), rd_o_cvd, rd_o_p0ref - ) - - mo_solve_nonhydro_stencil_67( - rho, - theta_v, - exner, - rd_o_cvd, - rd_o_p0ref, - offset_provider={}, - ) - - assert np.allclose(theta_v, theta_v_ref) - assert np.allclose(exner, exner_ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_68.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_68.py deleted file mode 100644 index 6465852303..0000000000 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_68.py +++ /dev/null @@ -1,82 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_68 import ( - mo_solve_nonhydro_stencil_68, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, random_mask -from .test_utils.simple_mesh import SimpleMesh - - -def mo_solve_nonhydro_stencil_68_numpy( - mask_prog_halo_c: np.array, - rho_now: np.array, - theta_v_now: np.array, - exner_new: np.array, - exner_now: np.array, - rho_new: np.array, - theta_v_new: np.array, - cvd_o_rd: float, -) -> np.array: - - mask_prog_halo_c = np.expand_dims(mask_prog_halo_c, axis=-1) - - theta_v_new = np.where( - mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) - / rho_new, - theta_v_new, - ) - return theta_v_new - - -def test_mo_solve_nonhydro_stencil_68(): - mesh = SimpleMesh() - - mask_prog_halo_c = random_mask(mesh, CellDim) - rho_now = random_field(mesh, CellDim, KDim) - theta_v_now = random_field(mesh, CellDim, KDim) - exner_new = random_field(mesh, CellDim, KDim) - exner_now = random_field(mesh, CellDim, KDim) - rho_new = random_field(mesh, CellDim, KDim) - theta_v_new = random_field(mesh, CellDim, KDim) - cvd_o_rd = 10.0 - - ref = mo_solve_nonhydro_stencil_68_numpy( - np.asarray(mask_prog_halo_c), - np.asarray(rho_now), - np.asarray(theta_v_now), - np.asarray(exner_new), - np.asarray(exner_now), - np.asarray(rho_new), - np.asarray(theta_v_new), - np.asarray(cvd_o_rd), - ) - mo_solve_nonhydro_stencil_68( - mask_prog_halo_c, - rho_now, - theta_v_now, - exner_new, - exner_now, - rho_new, - theta_v_new, - cvd_o_rd, - offset_provider={}, - ) - assert np.allclose(theta_v_new, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py deleted file mode 100644 index b14a2fda1b..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_01 import ( - mo_velocity_advection_stencil_01, -) -from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_01_numpy( - e2c2e: np.array, vn: np.array, rbf_vec_coeff_e: np.array -) -> np.array: - rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) - vt = np.sum(vn[e2c2e] * rbf_vec_coeff_e, axis=1) - return vt - - -def test_mo_velocity_advection_stencil_01(): - mesh = SimpleMesh() - - vn = random_field(mesh, EdgeDim, KDim) - rbf_vec_coeff = random_field(mesh, EdgeDim, E2C2EDim) - vt = zero_field(mesh, EdgeDim, KDim) - - ref = mo_velocity_advection_stencil_01_numpy( - mesh.e2c2e, - np.asarray(vn), - np.asarray(rbf_vec_coeff), - ) - mo_velocity_advection_stencil_01( - vn, - rbf_vec_coeff, - vt, - offset_provider={"E2C2E": mesh.get_e2c2e_offset_provider()}, - ) - assert np.allclose(vt, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_04.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_04.py deleted file mode 100644 index fbb9656ff7..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_04.py +++ /dev/null @@ -1,50 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_04 import ( - mo_velocity_advection_stencil_04, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_04_numpy( - vn: np.array, ddxn_z_full: np.array, ddxt_z_full: np.array, vt: np.array -) -> np.array: - z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full - return z_w_concorr_me - - -def test_mo_velocity_advection_stencil_04(): - mesh = SimpleMesh() - - vn = random_field(mesh, EdgeDim, KDim) - ddxn_z_full = random_field(mesh, EdgeDim, KDim) - ddxt_z_full = random_field(mesh, EdgeDim, KDim) - vt = random_field(mesh, EdgeDim, KDim) - z_w_concorr_me = zero_field(mesh, EdgeDim, KDim) - - ref = mo_velocity_advection_stencil_04_numpy( - np.asarray(vn), - np.asarray(ddxn_z_full), - np.asarray(ddxt_z_full), - np.asarray(vt), - ) - mo_velocity_advection_stencil_04( - vn, ddxn_z_full, ddxt_z_full, vt, z_w_concorr_me, offset_provider={} - ) - assert np.allclose(z_w_concorr_me, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_05.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_05.py deleted file mode 100644 index 79b08a2c77..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_05.py +++ /dev/null @@ -1,58 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_05 import ( - mo_velocity_advection_stencil_05, -) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_05_numpy( - vn: np.array, vt: np.array -) -> tuple[np.array]: - vn_ie = vn - z_vt_ie = vt - z_kin_hor_e = 0.5 * ((vn * vn) + (vt * vt)) - return vn_ie, z_vt_ie, z_kin_hor_e - - -def test_mo_velocity_advection_stencil_05(): - mesh = SimpleMesh() - - vn = random_field(mesh, EdgeDim, KDim) - vt = random_field(mesh, EdgeDim, KDim) - - vn_ie = zero_field(mesh, EdgeDim, KDim) - z_vt_ie = zero_field(mesh, EdgeDim, KDim) - z_kin_hor_e = zero_field(mesh, EdgeDim, KDim) - - vn_ie_ref, z_vt_ie_ref, z_kin_hor_e_ref = mo_velocity_advection_stencil_05_numpy( - np.asarray(vn), np.asarray(vt) - ) - - mo_velocity_advection_stencil_05( - vn, - vt, - vn_ie, - z_vt_ie, - z_kin_hor_e, - offset_provider={}, - ) - assert np.allclose(vn_ie, vn_ie_ref) - assert np.allclose(z_vt_ie, z_vt_ie_ref) - assert np.allclose(z_kin_hor_e, z_kin_hor_e_ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_07.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_07.py deleted file mode 100644 index 2a9a89816a..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_07.py +++ /dev/null @@ -1,91 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_07 import ( - mo_velocity_advection_stencil_07, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_07_numpy( - e2c: np.array, - e2v: np.array, - vn_ie: np.array, - inv_dual_edge_length: np.array, - w: np.array, - z_vt_ie: np.array, - inv_primal_edge_length: np.array, - tangent_orientation: np.array, - z_w_v: np.array, -) -> np.array: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - - w_e2c = w[e2c] - z_w_v_e2v = z_w_v[e2v] - - red_w = w_e2c[:, 0] - w_e2c[:, 1] - red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] - - z_v_grad_w = ( - vn_ie * inv_dual_edge_length * red_w - + z_vt_ie * inv_primal_edge_length * tangent_orientation * red_z_w_v - ) - return z_v_grad_w - - -def test_mo_velocity_advection_stencil_07(): - mesh = SimpleMesh() - - vn_ie = random_field(mesh, EdgeDim, KDim) - inv_dual_edge_length = random_field(mesh, EdgeDim) - w = random_field(mesh, CellDim, KDim) - z_vt_ie = random_field(mesh, EdgeDim, KDim) - inv_primal_edge_length = random_field(mesh, EdgeDim) - tangent_orientation = random_field(mesh, EdgeDim) - z_w_v = random_field(mesh, VertexDim, KDim) - z_v_grad_w = zero_field(mesh, EdgeDim, KDim) - - ref = mo_velocity_advection_stencil_07_numpy( - mesh.e2c, - mesh.e2v, - np.asarray(vn_ie), - np.asarray(inv_dual_edge_length), - np.asarray(w), - np.asarray(z_vt_ie), - np.asarray(inv_primal_edge_length), - np.asarray(tangent_orientation), - np.asarray(z_w_v), - ) - - mo_velocity_advection_stencil_07( - vn_ie, - inv_dual_edge_length, - w, - z_vt_ie, - inv_primal_edge_length, - tangent_orientation, - z_w_v, - z_v_grad_w, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - "E2V": mesh.get_e2v_offset_provider(), - }, - ) - assert np.allclose(z_v_grad_w, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py deleted file mode 100644 index 9c174d5572..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_08 import ( - mo_velocity_advection_stencil_08, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_08_numpy( - c2e: np.array, z_kin_hor_e: np.array, e_bln_c_s: np.array -) -> np.array: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - z_ekinh = np.sum(z_kin_hor_e[c2e] * e_bln_c_s, axis=1) - return z_ekinh - - -def test_mo_velocity_advection_stencil_08(): - mesh = SimpleMesh() - - z_kin_hor_e = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) - z_ekinh = zero_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_08_numpy( - mesh.c2e, - np.asarray(z_kin_hor_e), - np.asarray(e_bln_c_s), - ) - mo_velocity_advection_stencil_08( - z_kin_hor_e, - e_bln_c_s, - z_ekinh, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, - ) - assert np.allclose(z_ekinh, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py deleted file mode 100644 index 93639d8cd4..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py +++ /dev/null @@ -1,51 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_09 import ( - mo_velocity_advection_stencil_09, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_09_numpy( - c2e: np.array, z_w_concorr_me: np.array, e_bln_c_s: np.array -) -> np.array: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - z_w_concorr_mc = np.sum(z_w_concorr_me[c2e] * e_bln_c_s, axis=1) - return z_w_concorr_mc - - -def test_mo_velocity_advection_stencil_09(): - mesh = SimpleMesh() - - z_w_concorr_me = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) - z_w_concorr_mc = zero_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_09_numpy( - mesh.c2e, - np.asarray(z_w_concorr_me), - np.asarray(e_bln_c_s), - ) - mo_velocity_advection_stencil_09( - z_w_concorr_me, - e_bln_c_s, - z_w_concorr_mc, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, - ) - assert np.allclose(z_w_concorr_mc, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_11.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_11.py deleted file mode 100644 index 1b56752e0b..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_11.py +++ /dev/null @@ -1,39 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_11 import ( - mo_velocity_advection_stencil_11, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_11_numpy(w: np.array) -> np.array: - z_w_con_c = w - return z_w_con_c - - -def test_mo_velocity_advection_stencil_11(): - mesh = SimpleMesh() - - w = random_field(mesh, CellDim, KDim) - z_w_con_c = zero_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_11_numpy(w) - mo_velocity_advection_stencil_11(w, z_w_con_c, offset_provider={}) - - assert np.allclose(z_w_con_c, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_12.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_12.py deleted file mode 100644 index e6f27d1387..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_12.py +++ /dev/null @@ -1,40 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_12 import ( - mo_velocity_advection_stencil_12, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_12_numpy(z_w_con_c: np.array) -> np.array: - z_w_con_c = np.zeros_like(z_w_con_c) - return z_w_con_c - - -def test_mo_velocity_advection_stencil_12(): - mesh = SimpleMesh() - - z_w_con_c = random_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_12_numpy(np.asarray(z_w_con_c)) - mo_velocity_advection_stencil_12( - z_w_con_c, - offset_provider={}, - ) - assert np.allclose(z_w_con_c, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_13.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_13.py deleted file mode 100644 index 0d581df222..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_13.py +++ /dev/null @@ -1,47 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_13 import ( - mo_velocity_advection_stencil_13, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_13_numpy( - w_concorr_c: np.array, - z_w_con_c: np.array, -) -> np.array: - z_w_con_c = z_w_con_c - w_concorr_c - return z_w_con_c - - -def test_mo_velocity_advection_stencil_13(): - mesh = SimpleMesh() - - z_w_con_c = random_field(mesh, CellDim, KDim) - w_concorr_c = random_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_13_numpy( - np.asarray(w_concorr_c), np.asarray(z_w_con_c) - ) - mo_velocity_advection_stencil_13( - w_concorr_c, - z_w_con_c, - offset_provider={}, - ) - assert np.allclose(z_w_con_c, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_14.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_14.py deleted file mode 100644 index 748974ef3e..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_14.py +++ /dev/null @@ -1,90 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_14 import ( - mo_velocity_advection_stencil_14, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, random_mask, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_14_numpy( - ddqz_z_half: np.array, - z_w_con_c: np.array, - cfl_w_limit, - dtime, -) -> tuple[np.array]: - num_rows, num_cols = z_w_con_c.shape - cfl_clipping = np.where( - np.abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, - np.ones([num_rows, num_cols]), - np.zeros_like(z_w_con_c), - ) - num_rows, num_cols = cfl_clipping.shape - pre_levelmask = np.where( - cfl_clipping == 1.0, np.ones([num_rows, num_cols]), np.zeros_like(cfl_clipping) - ) - vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c - ) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c - ) - - return cfl_clipping, pre_levelmask, vcfl, z_w_con_c - - -def test_mo_velocity_advection_stencil_14(): - mesh = SimpleMesh() - - ddqz_z_half = random_field(mesh, CellDim, KDim) - z_w_con_c = random_field(mesh, CellDim, KDim) - cfl_clipping = random_mask(mesh, CellDim, KDim, dtype=bool) - pre_levelmask = random_mask( - mesh, CellDim, KDim, dtype=bool - ) # TODO should be just a K field - vcfl = zero_field(mesh, CellDim, KDim) - cfl_w_limit = 5.0 - dtime = 9.0 - - ( - cfl_clipping_ref, - pre_levelmask_ref, - vcfl_ref, - z_w_con_c_ref, - ) = mo_velocity_advection_stencil_14_numpy( - np.asarray(ddqz_z_half), - np.asarray(z_w_con_c), - cfl_w_limit, - dtime, - ) - - mo_velocity_advection_stencil_14( - ddqz_z_half, - z_w_con_c, - cfl_clipping, - pre_levelmask, - vcfl, - cfl_w_limit, - dtime, - offset_provider={}, - ) - assert np.allclose(cfl_clipping, cfl_clipping_ref) - assert np.allclose(pre_levelmask, pre_levelmask_ref) - assert np.allclose(vcfl, vcfl_ref) - assert np.allclose(z_w_con_c, z_w_con_c_ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_15.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_15.py deleted file mode 100644 index 31e7a3dc26..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_15.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_15 import ( - mo_velocity_advection_stencil_15, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_15_numpy( - z_w_con_c: np.array, -): - z_w_con_c_full = 0.5 * (z_w_con_c[:, :-1] + z_w_con_c[:, 1:]) - return z_w_con_c_full - - -def test_mo_velocity_advection_stencil_15(): - mesh = SimpleMesh() - - z_w_con_c = random_field(mesh, CellDim, KDim, extend={KDim: 1}) - - z_w_con_c_full = zero_field(mesh, CellDim, KDim) - - z_w_con_c_full_ref = mo_velocity_advection_stencil_15_numpy( - np.asarray(z_w_con_c), - ) - - mo_velocity_advection_stencil_15( - z_w_con_c, - z_w_con_c_full, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(z_w_con_c_full, z_w_con_c_full_ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py deleted file mode 100644 index 953d42e0f6..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py +++ /dev/null @@ -1,49 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_17 import ( - mo_velocity_advection_stencil_17, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_17_numpy( - c2e: np.array, e_bln_c_s: np.array, z_v_grad_w: np.array, ddt_w_adv: np.array -) -> np.array: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - ddt_w_adv = ddt_w_adv + np.sum(z_v_grad_w[c2e] * e_bln_c_s, axis=1) - return ddt_w_adv - - -def test_mo_velocity_advection_stencil_17(): - mesh = SimpleMesh() - - z_v_grad_w = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) - ddt_w_adv = random_field(mesh, CellDim, KDim) - - ref = mo_velocity_advection_stencil_17_numpy( - mesh.c2e, np.asarray(e_bln_c_s), np.asarray(z_v_grad_w), np.asarray(ddt_w_adv) - ) - mo_velocity_advection_stencil_17( - e_bln_c_s, - z_v_grad_w, - ddt_w_adv, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, - ) - assert np.allclose(ddt_w_adv, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py deleted file mode 100644 index b01db84bf6..0000000000 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py +++ /dev/null @@ -1,112 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_18 import ( - mo_velocity_advection_stencil_18, -) -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - -from .test_utils.helpers import random_field, random_mask -from .test_utils.simple_mesh import SimpleMesh - - -def mo_velocity_advection_stencil_18_numpy( - c2e2c0: np.array, - levelmask: np.array, - cfl_clipping: np.array, - owner_mask: np.array, - z_w_con_c: np.array, - ddqz_z_half: np.array, - area: np.array, - geofac_n2s: np.array, - w: np.array, - ddt_w_adv: np.array, - scalfac_exdiff: float, - cfl_w_limit: float, - dtime: float, -): - levelmask = np.expand_dims(levelmask, axis=0) - owner_mask = np.expand_dims(owner_mask, axis=-1) - area = np.expand_dims(area, axis=-1) - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - - difcoef = np.where( - (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(z_w_con_c) * dtime / ddqz_z_half - cfl_w_limit * dtime, - ), - 0, - ) - - ddt_w_adv = np.where( - (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), - ddt_w_adv + difcoef * area * np.sum(w[c2e2c0] * geofac_n2s, axis=1), - ddt_w_adv, - ) - - return ddt_w_adv - - -def test_mo_velocity_advection_stencil_18(): - mesh = SimpleMesh() - - levelmask = random_mask(mesh, KDim) - cfl_clipping = random_mask(mesh, CellDim, KDim) - owner_mask = random_mask(mesh, CellDim) - z_w_con_c = random_field(mesh, CellDim, KDim) - ddqz_z_half = random_field(mesh, CellDim, KDim) - area = random_field(mesh, CellDim) - geofac_n2s = random_field(mesh, CellDim, C2E2CODim) - w = random_field(mesh, CellDim, KDim) - ddt_w_adv = random_field(mesh, CellDim, KDim) - scalfac_exdiff = 10.0 - cfl_w_limit = 3.0 - dtime = 2.0 - - ref = mo_velocity_advection_stencil_18_numpy( - mesh.c2e2cO, - np.asarray(levelmask), - np.asarray(cfl_clipping), - np.asarray(owner_mask), - np.asarray(z_w_con_c), - np.asarray(ddqz_z_half), - np.asarray(area), - np.asarray(geofac_n2s), - np.asarray(w), - np.asarray(ddt_w_adv), - scalfac_exdiff, - cfl_w_limit, - dtime, - ) - - mo_velocity_advection_stencil_18( - levelmask, - cfl_clipping, - owner_mask, - z_w_con_c, - ddqz_z_half, - area, - geofac_n2s, - w, - ddt_w_adv, - scalfac_exdiff, - cfl_w_limit, - dtime, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, - ) - - assert np.allclose(ddt_w_adv, ref) diff --git a/atm_dyn_iconam/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/atm_dyn_iconam/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py deleted file mode 100644 index 6225933325..0000000000 --- a/atm_dyn_iconam/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ /dev/null @@ -1,58 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( - temporary_field_for_grid_point_cold_pools_enhancement, -) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def temporary_field_for_grid_point_cold_pools_enhancement_numpy( - c2e2c: np.array, theta_v: np.array, theta_ref_mc: np.array, thresh_tdiff -) -> np.array: - tdiff = theta_v - np.sum(theta_v[c2e2c], axis=1) / 3 - trefdiff = theta_ref_mc - np.sum(theta_ref_mc[c2e2c], axis=1) / 3 - - enh_diffu_3d = np.where( - ((tdiff - trefdiff) < thresh_tdiff) & (trefdiff < 0), - (thresh_tdiff - tdiff + trefdiff) * 5e-4, - -1.7976931348623157e308, - ) - - return enh_diffu_3d - - -def test_temporary_field_for_grid_point_cold_pools_enhancement(): - mesh = SimpleMesh() - - theta_v = random_field(mesh, CellDim, KDim) - theta_ref_mc = random_field(mesh, CellDim, KDim) - enh_diffu_3d = zero_field(mesh, CellDim, KDim) - thresh_tdiff = 5.0 - - ref = temporary_field_for_grid_point_cold_pools_enhancement_numpy( - mesh.c2e2c, np.asarray(theta_v), np.asarray(theta_ref_mc), thresh_tdiff - ) - temporary_field_for_grid_point_cold_pools_enhancement( - theta_v, - theta_ref_mc, - enh_diffu_3d, - thresh_tdiff, - offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, - ) - assert np.allclose(enh_diffu_3d, ref) diff --git a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py deleted file mode 100644 index 22b7ea4a4c..0000000000 --- a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ /dev/null @@ -1,80 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( - temporary_fields_for_turbulence_diagnostics, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -def mo_nh_diffusion_stencil_02_numpy( - c2e: np.array, - kh_smag_ec: np.array, - vn: np.array, - e_bln_c_s: np.array, - geofac_div: np.array, - diff_multfac_smag: np.array, -) -> tuple[np.array]: - geofac_div = np.expand_dims(geofac_div, axis=-1) - vn_geofac = vn[c2e] * geofac_div - div = np.sum(vn_geofac, axis=1) - - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - diff_multfac_smag = np.expand_dims(diff_multfac_smag, axis=0) - mul = kh_smag_ec[c2e] * e_bln_c_s - summed = np.sum(mul, axis=1) - kh_c = summed / diff_multfac_smag - - return div, kh_c - - -def test_mo_nh_diffusion_stencil_02(): - mesh = SimpleMesh() - - vn = random_field(mesh, EdgeDim, KDim) - geofac_div = random_field(mesh, CellDim, C2EDim) - kh_smag_ec = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) - diff_multfac_smag = random_field(mesh, KDim) - - kh_c = zero_field(mesh, CellDim, KDim) - div = zero_field(mesh, CellDim, KDim) - - div_ref, kh_c_ref = mo_nh_diffusion_stencil_02_numpy( - mesh.c2e, - np.asarray(kh_smag_ec), - np.asarray(vn), - np.asarray(e_bln_c_s), - np.asarray(geofac_div), - np.asarray(diff_multfac_smag), - ) - - temporary_fields_for_turbulence_diagnostics( - kh_smag_ec, - vn, - e_bln_c_s, - geofac_div, - diff_multfac_smag, - kh_c, - div, - offset_provider={ - "C2E": mesh.get_c2e_offset_provider(), - }, - ) - assert np.allclose(kh_c, kh_c_ref) - assert np.allclose(div, div_ref) diff --git a/atm_dyn_iconam/tests/test_update_theta_and_exner.py b/atm_dyn_iconam/tests/test_update_theta_and_exner.py deleted file mode 100644 index 942701ad8a..0000000000 --- a/atm_dyn_iconam/tests/test_update_theta_and_exner.py +++ /dev/null @@ -1,62 +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 - -import numpy as np - -from icon4py.atm_dyn_iconam.update_theta_and_exner import update_theta_and_exner -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh - - -def update_theta_and_exner_numpy( - z_temp: np.array, - area: np.array, - theta_v: np.array, - exner: np.array, - rd_o_cvd, -) -> tuple[np.array]: - area = np.expand_dims(area, axis=0) - z_theta = theta_v - theta_v = theta_v + (np.expand_dims(area, axis=-1) * z_temp) - exner = exner * (1.0 + rd_o_cvd * (theta_v / z_theta - 1.0)) - return theta_v, exner - - -def test_update_theta_and_exner(): - mesh = SimpleMesh() - - z_temp = random_field(mesh, CellDim, KDim) - area = random_field(mesh, CellDim) - theta_v = random_field(mesh, CellDim, KDim) - exner = random_field(mesh, CellDim, KDim) - rd_o_cvd = 5.0 - - theta_v_ref, exner_ref = update_theta_and_exner_numpy( - np.asarray(z_temp), - np.asarray(area), - np.asarray(theta_v), - np.asarray(exner), - rd_o_cvd, - ) - update_theta_and_exner( - z_temp, - area, - theta_v, - exner, - rd_o_cvd, - offset_provider={}, - ) - assert np.allclose(theta_v, theta_v_ref) - assert np.allclose(exner, exner_ref) diff --git a/atm_dyn_iconam/tests/test_utils/__init__.py b/atm_dyn_iconam/tests/test_utils/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/atm_dyn_iconam/tests/test_utils/__init__.py +++ /dev/null @@ -1,12 +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 diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 46f4ad777a..3c2f34e7af 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -11,11 +11,11 @@ flake8-debugger>=4.0.0 flake8-docstrings>=1.6.0 flake8-eradicate>=1.3.0 flake8-mutable>=1.2.0 -flake8-pyproject>=1.2.2 isort~=5.10 mypy>=0.942 pre-commit~=2.15 pytest>=6.1 +pytest-benchmark>=4.0.0 pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 diff --git a/common/README.md b/common/README.md deleted file mode 100644 index 2490a2c4a3..0000000000 --- a/common/README.md +++ /dev/null @@ -1,9 +0,0 @@ -# icon4py-common - -## Description - -Utilities shared by other ICON4Py components. - -## Installation instructions - -Check `README.md` file in the root of the repository. diff --git a/common/requirements-dev.txt b/common/requirements-dev.txt deleted file mode 100644 index 350767e37c..0000000000 --- a/common/requirements-dev.txt +++ /dev/null @@ -1,2 +0,0 @@ --r ../base-requirements-dev.txt --e . diff --git a/common/requirements.txt b/common/requirements.txt deleted file mode 100644 index 50242d5bde..0000000000 --- a/common/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ --r ../base-requirements.txt -. diff --git a/common/setup.cfg b/common/setup.cfg deleted file mode 100644 index 6cd2b35c10..0000000000 --- a/common/setup.cfg +++ /dev/null @@ -1,50 +0,0 @@ -# This file is mainly used to configure package creation with setuptools. -# Documentation: -# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files -# -[metadata] -name = icon4py_common -description = Icon inspired code in Python and GT4Py -long_description = file: README.md -long_description_content_type = text/markdown -url = https://github.com/C2SM/icon4py -author = ETH Zurich -author_email = gridtools@cscs.ch -license = gpl3 -license_files = LICENSE -platforms = Linux, Mac -classifiers = - Development Status :: 3 - Alpha - Intended Audience :: Science/Research - License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) - Operating System :: POSIX - Programming Language :: Python - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.10 - Programming Language :: Python :: Implementation :: CPython - Topic :: Scientific/Engineering :: Atmospheric Science - Topic :: Scientific/Engineering :: Mathematics - Topic :: Scientific/Engineering :: Physics -project_urls = - Source Code = https://github.com/C2SM/icon4py - -[options] -packages = find_namespace: -install_requires = - gt4py -python_requires = >=3.10 -package_dir = - = src -zip_safe = False - -[options.package_data] -# References: -# https://setuptools.pypa.io/en/latest/userguide/datafiles.html -# https://github.com/abravalheri/experiment-setuptools-package-data -* = *.md, *.rst, *.toml, *.txt, py.typed - -[options.packages.find] -where = src -exclude = - tests diff --git a/common/src/icon4py/common/__init__.py b/common/src/icon4py/common/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/common/src/icon4py/common/__init__.py +++ /dev/null @@ -1,12 +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 diff --git a/jenkins/spack-PR b/jenkins/spack-PR index 6782af678a..322dd30209 100644 --- a/jenkins/spack-PR +++ b/jenkins/spack-PR @@ -35,16 +35,6 @@ pipeline { """ } } - stage('Install with gt4py@stable') { - steps { - sh """ - . ./spack-c2sm/setup-env.sh - cd local_copy - spack env activate spack/gt4py-stable - spack install -v --test=root - """ - } - } } } } diff --git a/jenkins/spack-PR-icon b/jenkins/spack-PR-icon new file mode 100644 index 0000000000..7cadc494df --- /dev/null +++ b/jenkins/spack-PR-icon @@ -0,0 +1,95 @@ +def repo_identifier = "Project" +def icon_fork = "C2SM" +def icon_branch = "icon-dsl" +def icon_default = true + +def spack_fork = "C2SM" +def spack_branch = "main" +def spack_default = true + +def gt4py_fork = "GridTools" +def gt4py_branch = "icon4py_20230621" +def gt4py_default = true +String[] triggerPhrase = env.ghprbCommentBody.split(" ") + +def parseTriggerPhrase(it) { + fork = it.split("=")[1].split("/")[0] + branch = it.split("=")[1].split("/")[1] + return [fork, branch] +} +triggerPhrase.each { + if(it.contains("icon${repo_identifier}")) { + (icon_fork, icon_branch) = parseTriggerPhrase(it) + icon_default = false + } +} +triggerPhrase.each { + if(it.contains("spack${repo_identifier}")) { + (spack_fork, spack_branch) = parseTriggerPhrase(it) + spack_default = false + } +} +triggerPhrase.each { + if(it.contains("gt4py${repo_identifier}")) { + (gt4py_fork, gt4py_branch) = parseTriggerPhrase(it) + gt4py_default = false + } +} + +pipeline { + agent none + options{ + timeout(time: 3, unit: 'HOURS') + } + stages { + stage('Tests') { + matrix { + agent { label "${NODENAME}" } + axes { + axis { + name 'NODENAME' + values 'daint' + } + } + post { + always { + echo 'Cleaning up workspace' + deleteDir() + } + } + stages { + stage('Clone Repos') { + steps { + sh """ + git clone --depth 1 --recurse-submodules --shallow-submodules -b '${spack_branch}' https://github.com/${spack_fork}/spack-c2sm.git + git clone --depth 1 --recurse-submodules --shallow-submodules -b '${icon_branch}' git@github.com:${icon_fork}/icon-exclaim.git + git clone --depth 1 --recurse-submodules --shallow-submodules -b '${gt4py_branch}' git@github.com:${gt4py_fork}/gt4py.git + git clone . icon4py + """ + } + } + stage('Build with icon4py@main') { + steps { + sh """ + . ./spack-c2sm/setup-env.sh + cd icon-exclaim + spack env activate config/cscs/spack/latest/daint_dsl_dev + spack install -v --until build + """ + } + } + stage('Test with icon4py@main') { + steps { + sh """ + . ./spack-c2sm/setup-env.sh + cd icon-exclaim + spack env activate config/cscs/spack/latest/daint_dsl_dev + spack install -v --test=root + """ + } + } + } + } + } + } +} diff --git a/jenkins/spack-PR-stable b/jenkins/spack-PR-stable new file mode 100644 index 0000000000..78b85b9961 --- /dev/null +++ b/jenkins/spack-PR-stable @@ -0,0 +1,42 @@ +pipeline { + agent none + stages { + stage('Tests') { + matrix { + agent { label "${NODENAME}" } + axes { + axis { + name 'NODENAME' + values 'daint' + } + } + post { + always { + echo 'Cleaning up workspace' + deleteDir() + } + } + stages { + stage('Clone Repos') { + steps { + sh """ + git clone --depth 1 --recurse-submodules --shallow-submodules https://github.com/C2SM/spack-c2sm.git + git clone . local_copy + """ + } + } + stage('Install with gt4py@stable') { + steps { + sh """ + . ./spack-c2sm/setup-env.sh + cd local_copy + spack env activate spack/gt4py-stable + spack install -v --test=root + """ + } + } + } + } + } + } +} diff --git a/.flake8 b/model/.flake8 similarity index 91% rename from .flake8 rename to model/.flake8 index b37b750f96..31cecff5ab 100644 --- a/.flake8 +++ b/model/.flake8 @@ -40,6 +40,3 @@ rst-roles = py:attr, attr, py:exc, exc, py:obj, obj, - -per-file-ignores = - tools/src/icon4pytools/icon4pygen/icochainsize.py:E800 diff --git a/.license_header.txt b/model/.license_header.txt similarity index 100% rename from .license_header.txt rename to model/.license_header.txt diff --git a/.pre-commit-config.yaml b/model/.pre-commit-config.yaml similarity index 90% rename from .pre-commit-config.yaml rename to model/.pre-commit-config.yaml index 27e1a03f9e..04f4ab26a5 100644 --- a/.pre-commit-config.yaml +++ b/model/.pre-commit-config.yaml @@ -5,7 +5,7 @@ default_language_version: # Remove frozen version once we migrated away from tsa node: 17.9.1 minimum_pre_commit_version: 2.20.0 -exclude: "tools/.*" +exclude: "tools/.*|model/common/.*" repos: - repo: meta @@ -55,7 +55,7 @@ repos: --comment-style, "|#|", --license-filepath, - .license_header.txt, + model/.license_header.txt, --fuzzy-match-generates-todo, ] @@ -116,13 +116,19 @@ repos: - flake8-eradicate - flake8-mutable - pygments + args: + [ + --config=model/.flake8, + model/atmosphere/dycore/src/icon4py/, + model/common/src/icon4py/, + ] - repo: local hooks: - id: mypy name: mypy static type checker - entry: | - mypy --install-types --non-interactive -p icon4py.common + entry: bash -c 'echo mypy disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- language: system types_or: [python, pyi] always_run: true diff --git a/LICENSE b/model/LICENSE similarity index 100% rename from LICENSE rename to model/LICENSE diff --git a/model/README.md b/model/README.md new file mode 100644 index 0000000000..d591915055 --- /dev/null +++ b/model/README.md @@ -0,0 +1,24 @@ +# ICON4Py Model + +This folder contains Python implementations for multiple ICON components. + +It includes the following packages: + +- `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model +- `common`: Contains shared functionality that is required by multiple components. + +## Installation Instructions + +You can follow the general installation instructions used for the entire icon4py repository, or install each namespace package within this folder independently using the individual folder-specific `requirements.txt` or `requirements-dev.txt` files. + +In the following example it is assumed that you have already created and activated a virtual environment. + +```bash +# changing into the corresponding directory +cd model/atmosphere/dycore + +# installing a development version +pip install -r requirements-dev.txt +``` + +**Note**: For more information specific to each component, please refer to the README in their respective subfolders. diff --git a/model/atmosphere/dycore/.bumpversion.cfg b/model/atmosphere/dycore/.bumpversion.cfg new file mode 100644 index 0000000000..1cad57c0b4 --- /dev/null +++ b/model/atmosphere/dycore/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/atmosphere/dycore/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/atmosphere/dycore/README.md b/model/atmosphere/dycore/README.md new file mode 100644 index 0000000000..714bea17d4 --- /dev/null +++ b/model/atmosphere/dycore/README.md @@ -0,0 +1,9 @@ +# icon4py-atmosphere-dycore + +## Description + +Contains code ported from ICON `src/atm_dyn_iconam`, which is the dynamical core of the ICON model. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml new file mode 100644 index 0000000000..d4aee88f45 --- /dev/null +++ b/model/atmosphere/dycore/pyproject.toml @@ -0,0 +1,118 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel>=0.40.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "icon4py-atmosphere-dycore" +description = "ICON dynamical core." +readme = "README.md" +requires-python = ">=3.10" +license = {file = "LICENSE"} +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "icon4py_common>=0.0.5", +] +dynamic = ['version'] + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +parallel = true +branch = true +source_pkgs = ['dycore'] + +[tool.isort] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +multi_line_output = 3 +use_parentheses = true +include_trailing_comma = true +force_grid_wrap = 0 + +[tool.mypy] +install_types = true +non_interactive = true +exclude = [ + '^tests/*.py', +] +disallow_incomplete_defs = true +disallow_untyped_defs = true +ignore_missing_imports = false +implicit_reexport = true +warn_unused_configs = true +warn_unused_ignores = true +warn_redundant_casts = true +show_column_numbers = true +show_error_codes = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.atmosphere.dycore.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.atmosphere.dycore' = ['py.typed'] diff --git a/model/atmosphere/dycore/requirements-dev.txt b/model/atmosphere/dycore/requirements-dev.txt new file mode 100644 index 0000000000..32d385e08d --- /dev/null +++ b/model/atmosphere/dycore/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements-dev.txt +-e ../../common +-e . diff --git a/model/atmosphere/dycore/requirements.txt b/model/atmosphere/dycore/requirements.txt new file mode 100644 index 0000000000..79787ec137 --- /dev/null +++ b/model/atmosphere/dycore/requirements.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements.txt +../../common +. diff --git a/common/setup.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py similarity index 55% rename from common/setup.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py index 9c9f7b81c8..dab7089554 100644 --- a/common/setup.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py @@ -10,9 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final -from setuptools import setup +from packaging import version as pkg_version -if __name__ == "__main__": - setup() +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py index f607cfb262..a0e864646a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py @@ -14,15 +14,19 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.atm_dyn_iconam.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( _truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner -from icon4py.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.atmosphere.dycore.update_theta_and_exner import ( + _update_theta_and_exner, +) +from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py index 64f79b4d83..333c21116d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py @@ -14,17 +14,17 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( _apply_nabla2_and_nabla4_global_to_vn, ) -from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.atm_dyn_iconam.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.atm_dyn_iconam.calculate_nabla4 import _calculate_nabla4 -from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim +from icon4py.model.atmosphere.dycore.calculate_nabla4 import _calculate_nabla4 +from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @field_operator @@ -47,7 +47,6 @@ def _apply_diffusion_to_vn( start_2nd_nudge_line_idx_e: int32, limited_area: bool, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _calculate_nabla4( u_vert, v_vert, @@ -72,9 +71,7 @@ def _apply_diffusion_to_vn( vn, nudgezone_diff, ), - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v - ), + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) if limited_area else where( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 4019b4a083..4e8e1585f3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,15 +14,17 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( _calculate_horizontal_gradients_for_turbulence, ) -from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w -from icon4py.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import ( + _calculate_nabla2_for_w, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim @field_operator @@ -46,14 +48,11 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], ]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( int32(0) < vert_idx, - _calculate_horizontal_gradients_for_turbulence( - w_old, geofac_grg_x, geofac_grg_y - ), + _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) @@ -70,9 +69,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( & (vert_idx < nrdmax) & (interior_idx <= horz_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, area, z_nabla2_c - ), + _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), w, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py index 37eee89f66..313e9ac500 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -26,13 +27,11 @@ def _apply_nabla2_and_nabla4_global_to_vn( diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_and_nabla4_global_to_vn( area_edge: Field[[EdgeDim], float], kh_smag_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py index 21fd39cc40..c5872ec69e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py @@ -10,10 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, maximum -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -34,7 +35,7 @@ def _apply_nabla2_and_nabla4_to_vn( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_and_nabla4_to_vn( area_edge: Field[[EdgeDim], float], kh_smag_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py index 4225d2e5a1..c2f1074899 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py @@ -10,10 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -27,13 +28,11 @@ def _apply_nabla2_to_vn_in_lateral_boundary( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_vn_in_lateral_boundary( z_nabla2_e: Field[[EdgeDim, KDim], float], area_edge: Field[[EdgeDim], float], vn: Field[[EdgeDim, KDim], float], fac_bdydiff_v: float, ): - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn - ) + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py index e77fe2872c..f8d62fe062 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @field_operator @@ -31,7 +32,7 @@ def _apply_nabla2_to_w( return w -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_w( area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py index 07561e6180..42128fd8ad 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -28,13 +29,11 @@ def _apply_nabla2_to_w_in_upper_damping_layer( return w -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_w_in_upper_damping_layer( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w - ) + _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py index 4501c70214..e051aba455 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py @@ -14,13 +14,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, ) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim @field_operator diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostics_for_turbulence.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py index e276884113..1c686fe580 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -24,11 +25,12 @@ def _calculate_diagnostics_for_turbulence( wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: div_ic = wgtfac_c * div + (1.0 - wgtfac_c) * div(Koff[-1]) - hdef_ic = (wgtfac_c * kh_c + (1.0 - wgtfac_c) * kh_c(Koff[-1])) ** 2 + # TODO(magdalena): change exponent back to int (workaround for gt4py) + hdef_ic = (wgtfac_c * kh_c + (1.0 - wgtfac_c) * kh_c(Koff[-1])) ** 2.0 return div_ic, hdef_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_diagnostics_for_turbulence( div: Field[[CellDim, KDim], float], kh_c: Field[[CellDim, KDim], float], @@ -36,4 +38,4 @@ def calculate_diagnostics_for_turbulence( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic, hdef_ic)) + _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:])) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 98fef2843a..49eec5e7cf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -14,13 +14,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( _temporary_field_for_grid_point_cold_pools_enhancement, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim @field_operator @@ -33,9 +33,7 @@ def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( theta_v, theta_ref_mc, thresh_tdiff ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d - ) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d) return kh_smag_e diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py index 5ba95c84ae..5a0a98d830 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @field_operator @@ -28,7 +29,7 @@ def _calculate_horizontal_gradients_for_turbulence( return dwdx, dwdy -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_horizontal_gradients_for_turbulence( w: Field[[CellDim, KDim], float], geofac_grg_x: Field[[CellDim, C2E2CODim], float], @@ -36,6 +37,4 @@ def calculate_horizontal_gradients_for_turbulence( dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], ): - _calculate_horizontal_gradients_for_turbulence( - w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy) - ) + _calculate_horizontal_gradients_for_turbulence(w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py index 74becbc268..8df12abc3f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, maximum, minimum, sqrt +from gt4py.next.ffront.fbuiltins import Field, int32, maximum, minimum, sqrt -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( E2C2V, E2ECV, ECVDim, @@ -97,7 +98,7 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( kh_smag_2 = kh_smag_2 * kh_smag_2 kh_smag_e = diff_multfac_smag * sqrt(kh_smag_2 + kh_smag_1) - + # TODO(magdalena): change exponent back to int (workaround for gt4py) z_nabla2_e = ( ( ( @@ -110,8 +111,8 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( ) ) - 2.0 * vn - ) * (inv_primal_edge_length**2) - + ) * (inv_primal_edge_length**2.0) + # TODO(magdalena): change exponent back to int (workaround for gt4py) z_nabla2_e = z_nabla2_e + ( ( ( @@ -124,7 +125,7 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( ) ) - 2.0 * vn - ) * (inv_vert_vert_length**2) + ) * (inv_vert_vert_length**2.0) z_nabla2_e = 4.0 * z_nabla2_e @@ -135,7 +136,7 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( return kh_smag_e, kh_smag_ec, z_nabla2_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag: Field[[KDim], float], tangent_orientation: Field[[EdgeDim], float], @@ -153,10 +154,10 @@ def calculate_nabla2_and_smag_coefficients_for_vn( kh_smag_ec: Field[[EdgeDim, KDim], float], z_nabla2_e: Field[[EdgeDim, KDim], float], smag_offset: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py similarity index 80% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py index 7dba291f33..92bf9e30e7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py @@ -14,11 +14,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator @@ -41,6 +43,4 @@ def calculate_nabla2_for_theta( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _calculate_nabla2_for_theta( - kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp - ) + _calculate_nabla2_for_theta(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py index fefa04f4c1..cdba6c34f8 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @field_operator @@ -25,7 +26,7 @@ def _calculate_nabla2_for_w( return z_nabla2_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_z.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_z.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py index 549a317343..4ce096cf61 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_z.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -27,7 +28,7 @@ def _calculate_nabla2_for_z( return z_nabla2_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_for_z( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py index 10a8d31ee3..0420b78a0a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( C2CE, C2E, C2EDim, @@ -34,7 +35,7 @@ def _calculate_nabla2_of_theta( return z_temp -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_of_theta( z_nabla2_e: Field[[EdgeDim, KDim], float], geofac_div: Field[[CEDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla4.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py similarity index 92% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla4.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py index 2fe428609b..dede3cb396 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla4.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( E2C2V, E2ECV, ECVDim, @@ -48,14 +49,15 @@ def _calculate_nabla4( u_vert(E2C2V[3]) * primal_normal_vert_v1(E2ECV[3]) + v_vert(E2C2V[3]) * primal_normal_vert_v2(E2ECV[3]) ) + # TODO(magdalena): change exponent back to int (workaround for gt4py) z_nabla4_e2 = 4.0 * ( - (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length**2 - + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length**2 + (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length**2.0 + + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length**2.0 ) return z_nabla4_e2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla4( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 81% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index b7126aaea2..37430b5448 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, max_over, maximum -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim @field_operator @@ -26,11 +27,9 @@ def _enhance_diffusion_coefficient_for_grid_point_cold_pools( return kh_smag_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: Field[[EdgeDim, KDim], float], enh_diffu_3d: Field[[CellDim, KDim], float], ): - _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d, out=kh_smag_e - ) + _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py index b7161dc359..d2e875b515 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim @field_operator @@ -44,13 +45,11 @@ def _mo_advection_traj_btraj_compute_o1_dsl( p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( @@ -72,7 +71,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_advection_traj_btraj_compute_o1_dsl( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py similarity index 82% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 8036d7cbc8..295be5cab3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import V2C, CellDim, KDim, V2CDim, VertexDim +from icon4py.model.common.dimension import V2C, CellDim, KDim, V2CDim, VertexDim @field_operator @@ -26,12 +27,10 @@ def _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( return p_vert_out -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( p_cell_in: Field[[CellDim, KDim], float], c_intp: Field[[VertexDim, V2CDim], float], p_vert_out: Field[[VertexDim, KDim], float], ): - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in, c_intp, out=p_vert_out - ) + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(p_cell_in, c_intp, out=p_vert_out) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py index 8bc132e640..0a9bf2d689 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @field_operator @@ -28,7 +29,7 @@ def _mo_intp_rbf_rbf_vec_interpol_vertex( return p_u_out, p_v_out -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in: Field[[EdgeDim, KDim], float], ptr_coeff_1: Field[[VertexDim, V2EDim], float], @@ -36,6 +37,4 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], ): - _mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out) - ) + _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_divrot_rot_vertex_ri_dsl.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py index 6d7f2025c7..96014394da 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_divrot_rot_vertex_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @field_operator @@ -26,7 +27,7 @@ def _mo_math_divrot_rot_vertex_ri_dsl( return rot_vec -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_math_divrot_rot_vertex_ri_dsl( vec_e: Field[[EdgeDim, KDim], float], geofac_rot: Field[[VertexDim, V2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_gradients_grad_green_gauss_cell_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_gradients_grad_green_gauss_cell_dsl.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py index 7941e725e1..64349e1eb6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_math_gradients_grad_green_gauss_cell_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @field_operator @@ -36,7 +37,7 @@ def _mo_math_gradients_grad_green_gauss_cell_dsl( return p_grad_1_u, p_grad_1_v, p_grad_2_u, p_grad_2_v -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_math_gradients_grad_green_gauss_cell_dsl( p_grad_1_u: Field[[CellDim, KDim], float], p_grad_1_v: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_4th_order_divdamp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_4th_order_divdamp.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py index d67de6b2ef..b3920312e6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_4th_order_divdamp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -28,7 +29,7 @@ def _mo_solve_nonhydro_4th_order_divdamp( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_4th_order_divdamp( scal_divdamp: Field[[KDim], float], z_graddiv2_vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_01.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_01.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_01.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_01.py index 249438cc86..fe6012af9f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_01.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_01.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_01() -> ( return z_rth_pr_1, z_rth_pr_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_01( z_rth_pr_1: Field[[CellDim, KDim], float], z_rth_pr_2: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_02.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py index 15e8c63191..ca293f5a91 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py @@ -10,10 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,14 +24,12 @@ def _mo_solve_nonhydro_stencil_02( exner_ref_mc: Field[[CellDim, KDim], float], exner_pr: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_exner_ex_pr = (1.0 + exner_exfac) * ( - exner - exner_ref_mc - ) - exner_exfac * exner_pr + z_exner_ex_pr = (1.0 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return z_exner_ex_pr, exner_pr -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_02( exner_exfac: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_03.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_03.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_03.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_03.py index 1ec3eb5f6e..295b27defe 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_03.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_03.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,7 +24,7 @@ def _mo_solve_nonhydro_stencil_03() -> Field[[CellDim, KDim], float]: return z_exner_ex_pr -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_03( z_exner_ex_pr: Field[[CellDim, KDim], float], ): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_04.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_04.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_04.py index 51d5b4a8c1..9b634c6581 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_04.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -30,10 +31,10 @@ def _mo_solve_nonhydro_stencil_04( return z_exner_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_04( wgtfacq_c: Field[[CellDim, KDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], z_exner_ic: Field[[CellDim, KDim], float], ): - _mo_solve_nonhydro_stencil_04(wgtfacq_c, z_exner_ex_pr, out=z_exner_ic) + _mo_solve_nonhydro_stencil_04(wgtfacq_c, z_exner_ex_pr, out=z_exner_ic[:, 3:]) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_05.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py index e0be2d3d69..8c46ecb93e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_05( return z_exner_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_05( wgtfac_c: Field[[CellDim, KDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_06.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_06.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_06.py index f17829c1cb..8c1500a7a7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_06.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_06.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_06( return z_dexner_dz_c_1 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_06( z_exner_ic: Field[[CellDim, KDim], float], inv_ddqz_z_full: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_07.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py index 2e83d33f22..69089dea5b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -29,7 +30,7 @@ def _mo_solve_nonhydro_stencil_07( return z_rth_pr_1, z_rth_pr_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_07( rho: Field[[CellDim, KDim], float], rho_ref_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py similarity index 92% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_08.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py index 66baf27743..5cb85b1374 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_08.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -35,7 +36,7 @@ def _mo_solve_nonhydro_stencil_08( return rho_ic, z_rth_pr_1, z_rth_pr_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_08( wgtfac_c: Field[[CellDim, KDim], float], rho: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py similarity index 94% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_09.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py index 8f6bfe68d1..8ccdf40670 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_09.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -40,7 +41,7 @@ def _mo_solve_nonhydro_stencil_09( return z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_09( wgtfac_c: Field[[CellDim, KDim], float], z_rth_pr_2: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py similarity index 95% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_10.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index 4ef1c522fa..b709120573 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -42,9 +43,7 @@ def _mo_solve_nonhydro_stencil_10( ]: z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now(Koff[-1]) + wgt_nnew_rth * rho_var(Koff[-1]) - z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var( - Koff[-1] - ) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var(Koff[-1]) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( @@ -72,7 +71,7 @@ def _mo_solve_nonhydro_stencil_10( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_10( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_lower.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_lower.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_lower.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_lower.py index 9c1791d3ee..d8fbe662df 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_lower.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_lower.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -24,7 +25,7 @@ def _mo_solve_nonhydro_stencil_11_lower() -> Field[[CellDim, KDim], float]: return z_theta_v_pr_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_11_lower( z_theta_v_pr_ic: Field[[CellDim, KDim], float], ): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_upper.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py similarity index 92% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_upper.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py index 9c38802955..bfb20255c9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_11_upper.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -33,7 +34,7 @@ def _mo_solve_nonhydro_stencil_11_upper( return z_theta_v_pr_ic, theta_v_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_11_upper( wgtfacq_c: Field[[CellDim, KDim], float], z_rth_pr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_12.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_12.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_12.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_12.py index 8833fa9a82..978b59506c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_12.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_12.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -31,7 +32,7 @@ def _mo_solve_nonhydro_stencil_12( return z_dexner_dz_c_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_12( z_theta_v_pr_ic: Field[[CellDim, KDim], float], d2dexdz2_fac1_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_13.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_13.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py index 0708e644d4..2557b57479 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_13.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -29,7 +30,7 @@ def _mo_solve_nonhydro_stencil_13( return z_rth_pr_1, z_rth_pr_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_13( rho: Field[[CellDim, KDim], float], rho_ref_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_14.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_14.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_14.py index 11b8fa250d..f30e44624b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_14.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_14() -> ( return z_rho_e, z_theta_v_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_14( z_rho_e: Field[[EdgeDim, KDim], float], z_theta_v_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_15.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_15.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_15.py index 52e3c0e58e..853919bd72 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_15.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_15.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_15() -> ( return z_rho_e, z_theta_v_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_15( z_rho_e: Field[[EdgeDim, KDim], float], z_theta_v_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py similarity index 94% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index c5476150ef..eb9ac90b69 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -11,10 +11,18 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, where -from icon4py.common.dimension import E2C, E2EC, CellDim, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import ( + E2C, + E2EC, + CellDim, + ECDim, + EdgeDim, + KDim, +) @field_operator @@ -32,13 +40,11 @@ def _compute_btraj( lvn_pos = where(p_vn > 0.0, True, False) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( @@ -153,7 +159,7 @@ def _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1( return z_rho_e, z_theta_v_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_17.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py index 7bd486c681..4c628b66cf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -27,15 +28,12 @@ def _mo_solve_nonhydro_stencil_17( ) -> Field[[EdgeDim, KDim], float]: scalfac_dd3d = broadcast(scalfac_dd3d, (EdgeDim, KDim)) z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d - * scalfac_dd3d - * inv_dual_edge_length - * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) + hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) ) return z_graddiv_vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_17( hmask_dd3d: Field[[EdgeDim], float], scalfac_dd3d: Field[[KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py similarity index 82% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_18.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py index 4b4b5d8e4a..5046f22bfa 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -22,13 +23,11 @@ def _mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = inv_dual_edge_length * ( - z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0]) - ) + z_gradh_exner = inv_dual_edge_length * (z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0])) return z_gradh_exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_19.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_19.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_19.py index d3262436b9..c7760de829 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_19.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim @field_operator @@ -31,7 +32,7 @@ def _mo_solve_nonhydro_stencil_19( return z_gradh_exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_19( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py index 643c4c9b3e..0477ff1fcb 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py @@ -11,11 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType 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 ( +from icon4py.model.common.dimension import ( E2C, E2EC, CellDim, @@ -35,7 +36,6 @@ def _mo_solve_nonhydro_stencil_20( z_dexner_dz_c_1: Field[[CellDim, KDim], float], z_dexner_dz_c_2: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - 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]))) @@ -49,25 +49,19 @@ def _mo_solve_nonhydro_stencil_20( ( 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_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_dexner_dz_c1_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0])) ) ) return z_gradh_exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_20( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_21.py similarity index 92% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_21.py index 5e42edfe24..800898dc94 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_21.py @@ -11,11 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType 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 ( +from icon4py.model.common.dimension import ( E2C, E2EC, CellDim, @@ -36,7 +37,6 @@ def _mo_solve_nonhydro_stencil_21( inv_dual_edge_length: Field[[EdgeDim], float], grav_o_cpd: float, ) -> Field[[EdgeDim, KDim], float]: - theta_v_0 = theta_v(as_offset(Koff, ikoffset(E2EC[0]))) theta_v_1 = theta_v(as_offset(Koff, ikoffset(E2EC[1]))) @@ -55,19 +55,19 @@ def _mo_solve_nonhydro_stencil_21( 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]) - + # TODO(magdalena): change exponent back to int (workaround for gt4py) 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_theta_0 + z_theta_1) ** 2.0 ) return z_hydro_corr -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_21( theta_v: Field[[CellDim, KDim], float], ikoffset: Field[[ECDim, KDim], int32], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py similarity index 85% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_22.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py index b4572802ac..84fb050f6e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_22.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, where -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -24,13 +25,11 @@ def _mo_solve_nonhydro_stencil_22( z_hydro_corr: Field[[EdgeDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = where( - ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner - ) + z_gradh_exner = where(ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner) return z_gradh_exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_22( ipeidx_dsl: Field[[EdgeDim, KDim], bool], pg_exdist: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_23.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_23.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_23.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_23.py index 5249a10cac..3f937bf53f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_23.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_23.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -39,7 +40,7 @@ def _mo_solve_nonhydro_stencil_23( return vn_nnew -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_23( vn_nnow: Field[[EdgeDim, KDim], float], ddt_vn_adv_ntl1: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_24.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py index d0ef535a51..a23544d444 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_24.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -27,13 +28,11 @@ def _mo_solve_nonhydro_stencil_24( dtime: float, cpd: float, ) -> Field[[EdgeDim, KDim], float]: - vn_nnew = vn_nnow + dtime * ( - ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner - ) + vn_nnew = vn_nnow + dtime * (ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner) return vn_nnew -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_24( vn_nnow: Field[[EdgeDim, KDim], float], ddt_vn_adv_ntl1: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_25.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_25.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_25.py index 2f30b08c2a..d523419189 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_25.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_25.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_25( return z_graddiv2_vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_25( geofac_grdiv: Field[[EdgeDim, E2C2EODim], float], z_graddiv_vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_26.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_26.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_26.py index bf0560df03..0083d04f64 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_26.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_26.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -27,7 +28,7 @@ def _mo_solve_nonhydro_stencil_26( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_26( z_graddiv_vn: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_27.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_27.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_27.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_27.py index 51a825e027..8f39d33ac9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_27.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_27.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -31,7 +32,7 @@ def _mo_solve_nonhydro_stencil_27( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_27( scal_divdamp: Field[[KDim], float], bdy_divdamp: Field[[KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_28.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_28.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_28.py index f34d308640..2fe1ee4f09 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_28.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_28.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -27,7 +28,7 @@ def _mo_solve_nonhydro_stencil_28( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_28( vn_incr: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_29.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_29.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_29.py index 5a6d21c467..5efe92475b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_29.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_29.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -27,7 +28,7 @@ def _mo_solve_nonhydro_stencil_29( return vn_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_29( grf_tend_vn: Field[[EdgeDim, KDim], float], vn_now: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_30.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_30.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py index 23a7e56332..e8016b5926 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_30.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( E2C2E, E2C2EO, E2C2EDim, @@ -41,7 +42,7 @@ def _mo_solve_nonhydro_stencil_30( return z_vn_avg, z_graddiv_vn, vt -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_30( e_flx_avg: Field[[EdgeDim, E2C2EODim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_31.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_31.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_31.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_31.py index 57e8c9a6ee..550b60121d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_31.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_31.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_31( return z_vn_avg -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_31( e_flx_avg: Field[[EdgeDim, E2C2EODim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_32.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_32.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_32.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_32.py index 65a1d89b4b..b81c00bcaa 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_32.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_32.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -29,7 +30,7 @@ def _mo_solve_nonhydro_stencil_32( return mass_fl_e, z_theta_v_fl_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_32( z_rho_e: Field[[EdgeDim, KDim], float], z_vn_avg: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_33.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_33.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_33.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_33.py index 9aef58b809..e096127694 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_33.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_33.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_33() -> ( return vn_traj, mass_flx_me -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_33( vn_traj: Field[[EdgeDim, KDim], float], mass_flx_me: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_34.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_34.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_34.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_34.py index ec833a5ad1..1ab1fab30d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_34.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_34.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -30,7 +31,7 @@ def _mo_solve_nonhydro_stencil_34( return vn_traj, mass_flx_me -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_34( z_vn_avg: Field[[EdgeDim, KDim], float], mass_fl_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_35.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_35.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py index 1fd43a65ac..d1f56f80cc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_35.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -28,7 +29,7 @@ def _mo_solve_nonhydro_stencil_35( return z_w_concorr_me -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_35( vn: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py similarity index 85% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_36.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py index 14e6a5eced..89299c0309 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_36.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.common.dimension import EdgeDim, KDim, Koff @field_operator @@ -29,11 +30,12 @@ def _mo_solve_nonhydro_stencil_36( ]: vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn(Koff[-1]) z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt(Koff[-1]) - z_kin_hor_e = 0.5 * (vn**2 + vt**2) + # TODO(magdalena): change exponent back to int (workaround for gt4py) + z_kin_hor_e = 0.5 * (vn**2.0 + vt**2.0) return vn_ie, z_vt_ie, z_kin_hor_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_36( wgtfac_e: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_37.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_37.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_37.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_37.py index 9017027c89..f4361b03b3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_37.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_37.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -32,7 +33,7 @@ def _mo_solve_nonhydro_stencil_37( return vn_ie, z_vt_ie, z_kin_hor_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_37( vn: Field[[EdgeDim, KDim], float], vt: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_38.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py index 03ecc80ef0..7d49922aaf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_38.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.common.dimension import EdgeDim, KDim, Koff @field_operator @@ -30,7 +31,7 @@ def _mo_solve_nonhydro_stencil_38( return vn_ie -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_38( vn: Field[[EdgeDim, KDim], float], wgtfacq_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_39.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index f7cdcf6cbc..84abeee0e7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -11,10 +11,18 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff +from icon4py.model.common.dimension import ( + C2E, + C2EDim, + CellDim, + EdgeDim, + KDim, + Koff, +) @field_operator @@ -24,15 +32,13 @@ def _mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: z_w_concorr_me_offset_1 = z_w_concorr_me(Koff[-1]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) z_w_concorr_mc_m0 = neighbor_sum(e_bln_c_s * z_w_concorr_me(C2E), axis=C2EDim) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return w_concorr_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_39( e_bln_c_s: Field[[CellDim, C2EDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py similarity index 76% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_40.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 0c859620c4..6c40169603 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -11,10 +11,18 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff +from icon4py.model.common.dimension import ( + C2E, + C2EDim, + CellDim, + EdgeDim, + KDim, + Koff, +) @field_operator @@ -27,15 +35,9 @@ def _mo_solve_nonhydro_stencil_40( z_w_concorr_me_offset_2 = z_w_concorr_me(Koff[-2]) z_w_concorr_me_offset_3 = z_w_concorr_me(Koff[-3]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) - z_w_concorr_mc_m2 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim - ) - z_w_concorr_mc_m3 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) + z_w_concorr_mc_m2 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim) + z_w_concorr_mc_m3 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim) return ( wgtfacq_c(Koff[-1]) * z_w_concorr_mc_m1 @@ -44,7 +46,7 @@ def _mo_solve_nonhydro_stencil_40( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_40( e_bln_c_s: Field[[CellDim, C2EDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_41.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_41.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_41.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_41.py index f2b68e6eb2..594d408f88 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_41.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_41.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -28,7 +29,7 @@ def _mo_solve_nonhydro_stencil_41( return z_flxdiv_mass, z_flxdiv_theta -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_41( geofac_div: Field[[CellDim, C2EDim], float], mass_fl_e: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_42.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py index dccc14a589..5b6b0d51e2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -32,15 +33,13 @@ def _mo_solve_nonhydro_stencil_42( cpd: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c ) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) return z_w_expl, z_contr_w_fl_l -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_42( z_w_expl: Field[[CellDim, KDim], float], w_nnow: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_43.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_43.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_43.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_43.py index dbee44360c..931dbb33e8 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_43.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_43.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -33,7 +34,7 @@ def _mo_solve_nonhydro_stencil_43( return z_w_expl, z_contr_w_fl_l -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_43( z_w_expl: Field[[CellDim, KDim], float], w_nnow: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_44.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_44.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_44.py index d039b1a889..2f5513d1c1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_44.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_44.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -35,7 +36,7 @@ def _mo_solve_nonhydro_stencil_44( return z_beta, z_alpha -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_44( z_beta: Field[[CellDim, KDim], float], exner_nnow: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45.py index 924b46544d..1b29686595 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,6 +24,6 @@ def _mo_solve_nonhydro_stencil_45() -> Field[[CellDim, KDim], float]: return z_alpha -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_45(z_alpha: Field[[CellDim, KDim], float]): _mo_solve_nonhydro_stencil_45(out=z_alpha) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45_b.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45_b.py similarity index 80% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45_b.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45_b.py index c0c6fb8604..ad2a58dc16 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_45_b.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_45_b.py @@ -11,18 +11,19 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator -def _mo_solve_nonhydro_stencil_45_b() -> (Field[[CellDim, KDim], float]): +def _mo_solve_nonhydro_stencil_45_b() -> Field[[CellDim, KDim], float]: return broadcast(0.0, (CellDim, KDim)) -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_45_b( z_q: Field[[CellDim, KDim], float], ): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_46.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_46.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_46.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_46.py index 5b2a7dc769..0367e92e92 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_46.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_46.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_46() -> ( return w_nnew, z_contr_w_fl_l -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_46( w_nnew: Field[[CellDim, KDim], float], z_contr_w_fl_l: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_47.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_47.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_47.py index f46e16d58b..2828c8c980 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_47.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_47.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_solve_nonhydro_stencil_47( return w_nnew, z_contr_w_fl_l -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_47( w_nnew: Field[[CellDim, KDim], float], z_contr_w_fl_l: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_48.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_48.py similarity index 94% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_48.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_48.py index 9bb2fb9205..3dc035c06a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_48.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_48.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -47,7 +48,7 @@ def _mo_solve_nonhydro_stencil_48( return (z_rho_expl, z_exner_expl) -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_48( z_rho_expl: Field[[CellDim, KDim], float], z_exner_expl: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_49.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_49.py similarity index 94% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_49.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_49.py index 67829f58b4..f46d37876a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_49.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_49.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -46,7 +47,7 @@ def _mo_solve_nonhydro_stencil_49( return z_rho_expl, z_exner_expl -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_49( z_rho_expl: Field[[CellDim, KDim], float], z_exner_expl: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_50.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_50.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py index 5300ae804e..b857cfbfc7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_50.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -30,7 +31,7 @@ def _mo_solve_nonhydro_stencil_50( return z_rho_expl, z_exner_expl -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_50( z_rho_expl: Field[[CellDim, KDim], float], z_exner_expl: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_51.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py index 78862390fc..7396f7f1ac 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -39,7 +40,7 @@ def _mo_solve_nonhydro_stencil_51( return z_q, w_nnew -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_51( z_q: Field[[CellDim, KDim], float], w_nnew: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py similarity index 95% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_52.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py index 29e4de0457..027cb0e473 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program, scan_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @scan_operator(axis=KDim, forward=True, init=(0.0, 0.0, True)) @@ -61,7 +62,7 @@ def _mo_solve_nonhydro_stencil_52( return z_q_res, w_res -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_52( vwind_impl_wgt: Field[[CellDim], float], theta_v_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_53.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_53.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py index 116b38f70d..c6b208980a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_53.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import program, scan_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @scan_operator(axis=KDim, forward=False, init=0.0) @@ -22,7 +23,7 @@ def _mo_solve_nonhydro_stencil_53_scan(w_state: float, z_q: float, w: float) -> return w + w_state * z_q -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_53( z_q: Field[[CellDim, KDim], float], w: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_54.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py index bc1f8eb1c1..d1efc7b3f6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -28,7 +29,7 @@ def _mo_solve_nonhydro_stencil_54( return w -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_54( z_raylfac: Field[[KDim], float], w_1: Field[[CellDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_55.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py index 32966a1026..a354730e7d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -41,21 +42,12 @@ def _mo_solve_nonhydro_stencil_55( rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( rho_ic * w - rho_ic(Koff[1]) * w(Koff[1]) ) - exner_new = ( - z_exner_expl - + exner_ref_mc - - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) - ) - theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new - ) + exner_new = z_exner_expl + exner_ref_mc - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) + theta_v_new = rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new return rho_new, exner_new, theta_v_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_55( z_rho_expl: Field[[CellDim, KDim], float], vwind_impl_wgt: Field[[CellDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_56_63.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py index 1b2994aed1..e5e70e7ba2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -23,13 +24,11 @@ def _mo_solve_nonhydro_stencil_56_63( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - z_dwdz_dd = inv_ddqz_z_full * ( - (w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1])) - ) + z_dwdz_dd = inv_ddqz_z_full * ((w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1]))) return z_dwdz_dd -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_56_63( inv_ddqz_z_full: Field[[CellDim, KDim], float], w: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_57.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_57.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_57.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_57.py index 42996f0e00..2257403e2f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_57.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_57.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,6 +24,6 @@ def _mo_solve_nonhydro_stencil_57() -> Field[[CellDim, KDim], float]: return mass_flx_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_57(mass_flx_ic: Field[[CellDim, KDim], float]): _mo_solve_nonhydro_stencil_57(out=mass_flx_ic) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_58.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py index f9f6dce659..17879f2dbe 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -26,13 +27,11 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], float], r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return mass_flx_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_58( z_contr_w_fl_l: Field[[CellDim, KDim], float], rho_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_59.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_59.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_59.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_59.py index 08a2dcc029..5818ca40d8 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_59.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_59.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -25,7 +26,7 @@ def _mo_solve_nonhydro_stencil_59( return exner_dyn_incr -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_59( exner: Field[[CellDim, KDim], float], exner_dyn_incr: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py similarity index 85% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_60.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py index 6419b51cfe..995322d72d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -25,13 +26,11 @@ def _mo_solve_nonhydro_stencil_60( ndyn_substeps_var: float, dtime: float, ) -> Field[[CellDim, KDim], float]: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) return exner_dyn_incr -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_60( exner: Field[[CellDim, KDim], float], ddt_exner_phy: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_61.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py index f2c159af52..2eec2993b4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_61.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -37,7 +38,7 @@ def _mo_solve_nonhydro_stencil_61( return rho_new, exner_new, w_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_61( rho_now: Field[[CellDim, KDim], float], grf_tend_rho: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_62.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py index dd8d47c135..551690b8f0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -27,7 +28,7 @@ def _mo_solve_nonhydro_stencil_62( return w_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_62( w_now: Field[[CellDim, KDim], float], grf_tend_w: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_64.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_64.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_64.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_64.py index ed1a54a86f..b929fc3d63 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_64.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_64.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,7 +24,7 @@ def _mo_solve_nonhydro_stencil_64() -> Field[[CellDim, KDim], float]: return mass_flx_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_64( mass_flx_ic: Field[[CellDim, KDim], float], ): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_65.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py index a4a5b36dd6..eea826badc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -29,14 +30,12 @@ def _mo_solve_nonhydro_stencil_65( r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return mass_flx_ic -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_65( rho_ic: Field[[CellDim, KDim], float], vwind_expl_wgt: Field[[CellDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_66.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_66.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py index e3dd3689ee..2159d74657 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_66.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, exp, log, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -31,7 +32,7 @@ def _mo_solve_nonhydro_stencil_66( return theta_v, exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_66( bdy_halo_c: Field[[CellDim], bool], rho: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_67.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py index 8b0ae3413d..6f46f314d2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, exp, log -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -30,7 +31,7 @@ def _mo_solve_nonhydro_stencil_67( return theta_v, exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_67( rho: Field[[CellDim, KDim], float], theta_v: Field[[CellDim, KDim], float], @@ -38,6 +39,4 @@ def mo_solve_nonhydro_stencil_67( rd_o_cvd: float, rd_o_p0ref: float, ): - _mo_solve_nonhydro_stencil_67( - rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner) - ) + _mo_solve_nonhydro_stencil_67(rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_68.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py index c06fac2a5f..15230b8afb 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -30,16 +31,13 @@ def _mo_solve_nonhydro_stencil_68( ) -> Field[[CellDim, KDim], float]: theta_v_new = where( mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new, + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new, theta_v_new, ) return theta_v_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_68( mask_prog_halo_c: Field[[CellDim], bool], rho_now: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_01.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_01.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_01.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_01.py index 1a7b5abe41..99817e31b6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_01.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_01.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_velocity_advection_stencil_01( return vt -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_01( vn: Field[[EdgeDim, KDim], float], rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_02.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py index e919ba2cfe..15a16dc312 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.common.dimension import EdgeDim, KDim, Koff @field_operator @@ -29,7 +30,7 @@ def _mo_velocity_advection_stencil_02( return vn_ie, z_kin_hor_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_02( wgtfac_e: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_03.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py index 4005a047e3..12a7528853 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.common.dimension import EdgeDim, KDim, Koff @field_operator @@ -27,7 +28,7 @@ def _mo_velocity_advection_stencil_03( return z_vt_ie -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_03( wgtfac_e: Field[[EdgeDim, KDim], float], vt: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_04.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py index c5f07bd115..c7c051da02 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -28,7 +29,7 @@ def _mo_velocity_advection_stencil_04( return z_w_concorr_me -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_04( vn: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float], @@ -36,6 +37,4 @@ def mo_velocity_advection_stencil_04( vt: Field[[EdgeDim, KDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_04( - vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me - ) + _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_05.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_05.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_05.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_05.py index 432a11c88e..8bf99fc4f0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_05.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_05.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -32,7 +33,7 @@ def _mo_velocity_advection_stencil_05( return vn_ie, z_vt_ie, z_kin_hor_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_05( vn: Field[[EdgeDim, KDim], float], vt: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_06.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py index 8c5ff54863..4b158fbd90 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.common.dimension import EdgeDim, KDim, Koff @field_operator @@ -31,7 +32,7 @@ def _mo_velocity_advection_stencil_06( return vn_ie -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_06( wgtfacq_e: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_07.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py index 105169cd4d..a08de19b3f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py @@ -11,10 +11,18 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.dimension import ( + E2C, + E2V, + CellDim, + EdgeDim, + KDim, + VertexDim, +) @field_operator @@ -29,12 +37,10 @@ def _mo_velocity_advection_stencil_07( ) -> Field[[EdgeDim, KDim], float]: return vn_ie * inv_dual_edge_length * ( w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * ( - z_w_v(E2V[0]) - z_w_v(E2V[1]) - ) + ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * (z_w_v(E2V[0]) - z_w_v(E2V[1])) -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_07( vn_ie: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_08.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py index 03231c3ac0..a76d719136 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_velocity_advection_stencil_08( return z_ekinh -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_08( z_kin_hor_e: Field[[EdgeDim, KDim], float], e_bln_c_s: Field[[CellDim, C2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_09.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py index 6c635c2c7c..7e6218b151 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_velocity_advection_stencil_09( return z_w_concorr_mc -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_09( z_w_concorr_me: Field[[EdgeDim, KDim], float], e_bln_c_s: Field[[CellDim, C2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_10.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 8b272d547f..4dd324603b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -22,14 +23,12 @@ def _mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc( - Koff[-1] - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc(Koff[-1]) return w_concorr_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_11.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_11.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_11.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_11.py index 6f5dcf2e80..01f0eb66a5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_11.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_11.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -25,7 +26,7 @@ def _mo_velocity_advection_stencil_11( return z_w_con_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_11( w: Field[[CellDim, KDim], float], z_w_con_c: Field[[CellDim, KDim], float] ): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_12.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_12.py similarity index 87% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_12.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_12.py index 182f19e332..92bceb53a6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_12.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_12.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -23,6 +24,6 @@ def _mo_velocity_advection_stencil_12() -> Field[[CellDim, KDim], float]: return z_w_con_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_12(z_w_con_c: Field[[CellDim, KDim], float]): _mo_velocity_advection_stencil_12(out=z_w_con_c) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_13.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_13.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_13.py index 91219a28cb..7b33d5bc59 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_13.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -26,7 +27,7 @@ def _mo_velocity_advection_stencil_13( return z_w_con_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_13( w_concorr_c: Field[[CellDim, KDim], float], z_w_con_c: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py similarity index 86% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_14.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index 0f7026bf66..3bce7d7ef2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, abs, broadcast, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -42,18 +43,14 @@ def _mo_velocity_advection_stencil_14( vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = where( - (cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c) - z_w_con_c = where( - (cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c) return cfl_clipping, pre_levelmask, vcfl, z_w_con_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_14( ddqz_z_half: Field[[CellDim, KDim], float], z_w_con_c: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_15.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_15.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_15.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_15.py index 6cb41848aa..6c29da8add 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_15.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_15.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -25,7 +26,7 @@ def _mo_velocity_advection_stencil_15( return z_w_con_c_full -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_15( z_w_con_c: Field[[CellDim, KDim], float], z_w_con_c_full: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py similarity index 78% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_16.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py index 694acef6b1..a1823ac6be 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator @@ -25,14 +26,12 @@ def _mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: ddt_w_adv = -z_w_con_c * ( - w(Koff[-1]) * coeff1_dwdz - - w(Koff[1]) * coeff2_dwdz - + w * (coeff2_dwdz - coeff1_dwdz) + w(Koff[-1]) * coeff1_dwdz - w(Koff[1]) * coeff2_dwdz + w * (coeff2_dwdz - coeff1_dwdz) ) return ddt_w_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_16( z_w_con_c: Field[[CellDim, KDim], float], w: Field[[CellDim, KDim], float], @@ -40,6 +39,4 @@ def mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ddt_w_adv: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_16( - z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:] - ) + _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:]) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_17.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_17.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_17.py index a63c86a73b..654b272a9e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_17.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -27,7 +28,7 @@ def _mo_velocity_advection_stencil_17( return ddt_w_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_17( e_bln_c_s: Field[[CellDim, C2EDim], float], z_v_grad_w: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_18.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py index 8c07662513..c4e8d32d2b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, abs, minimum, neighbor_sum, where -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @field_operator @@ -44,15 +45,14 @@ def _mo_velocity_advection_stencil_18( ddt_w_adv = where( levelmask & cfl_clipping & owner_mask, - ddt_w_adv - + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), + ddt_w_adv + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), ddt_w_adv, ) return ddt_w_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_18( levelmask: Field[[KDim], bool], cfl_clipping: Field[[CellDim, KDim], bool], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_19.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py index b1014009ef..621e6befed 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( E2C, E2EC, E2V, @@ -44,10 +45,7 @@ def _mo_velocity_advection_stencil_19( ): ddt_vn_adv = -( (coeff_gradekin(E2EC[0]) - coeff_gradekin(E2EC[1])) * z_kin_hor_e - + ( - -coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) - + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1]) - ) + + (-coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1])) + vt * (f_e + 0.5 * neighbor_sum(zeta(E2V), axis=E2VDim)) + neighbor_sum(z_w_con_c_full(E2C) * c_lin_e, axis=E2CDim) * (vn_ie - vn_ie(Koff[1])) @@ -56,7 +54,7 @@ def _mo_velocity_advection_stencil_19( return ddt_vn_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_19( z_kin_hor_e: Field[[EdgeDim, KDim], float], coeff_gradekin: Field[[ECDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_20.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py index d85ffe919a..173bd2e9b1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import ( Field, @@ -22,7 +23,7 @@ where, ) -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( E2C, E2C2EO, E2V, @@ -75,15 +76,13 @@ def _mo_velocity_advection_stencil_20( (levelmask | levelmask(Koff[1])) & (abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), ddt_vn_adv + difcoef * area_edge * neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim) - + tangent_orientation - * inv_primal_edge_length - * neighbor_sum(zeta(E2V), axis=E2VDim), + + tangent_orientation * inv_primal_edge_length * neighbor_sum(zeta(E2V), axis=E2VDim), ddt_vn_adv, ) return ddt_vn_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def mo_velocity_advection_stencil_20( levelmask: Field[[KDim], bool], c_lin_e: Field[[EdgeDim, E2CDim], float], diff --git a/common/src/icon4py/common/py.typed b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/py.typed similarity index 100% rename from common/src/icon4py/common/py.typed rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/py.typed diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py index d0b14550be..81f841143c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum, where -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim +from icon4py.model.common.dimension import C2E2C, C2E2CDim, CellDim, KDim @field_operator @@ -33,7 +34,7 @@ def _temporary_field_for_grid_point_cold_pools_enhancement( return enh_diffu_3d -@program +@program(grid_type=GridType.UNSTRUCTURED) def temporary_field_for_grid_point_cold_pools_enhancement( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py similarity index 91% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py index c0b536319f..a90aaddea5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -30,7 +31,7 @@ def _temporary_fields_for_turbulence_diagnostics( return kh_c, div -@program +@program(grid_type=GridType.UNSTRUCTURED) def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 84% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index ab0574e40e..e331a620cd 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,7 +15,14 @@ from gt4py.next.ffront.experimental import as_offset from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim, Koff +from icon4py.model.common.dimension import ( + C2CEC, + C2E2C, + CECDim, + CellDim, + KDim, + Koff, +) @field_operator @@ -29,7 +36,6 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( 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]))) @@ -40,20 +46,11 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( 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]) - ) + * (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]) - ) + * (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]) - ) + * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) ) z_temp = where( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py similarity index 85% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py index 7b44b0e3f7..bc53f68979 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator @@ -31,7 +32,7 @@ def _update_theta_and_exner( return theta_v, exner -@program +@program(grid_type=GridType.UNSTRUCTURED) def update_theta_and_exner( z_temp: Field[[CellDim, KDim], float], area: Field[[CellDim], float], @@ -39,6 +40,4 @@ def update_theta_and_exner( exner: Field[[CellDim, KDim], float], rd_o_cvd: float, ): - _update_theta_and_exner( - z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner) - ) + _update_theta_and_exner(z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner)) diff --git a/atm_dyn_iconam/__init__.py b/model/atmosphere/dycore/tests/__init__.py similarity index 100% rename from atm_dyn_iconam/__init__.py rename to model/atmosphere/dycore/tests/__init__.py diff --git a/atm_dyn_iconam/tests/__init__.py b/model/atmosphere/dycore/tests/conftest.py similarity index 54% rename from atm_dyn_iconam/tests/__init__.py rename to model/atmosphere/dycore/tests/conftest.py index 15dfdb0098..7ee1151470 100644 --- a/atm_dyn_iconam/tests/__init__.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -10,3 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pytest +from gt4py.next.program_processors.runners.roundtrip import executor + +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py new file mode 100644 index 0000000000..d9d058adf3 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -0,0 +1,51 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( + apply_nabla2_and_nabla4_global_to_vn, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestApplyNabla2AndNabla4GlobalToVn(StencilTest): + PROGRAM = apply_nabla2_and_nabla4_global_to_vn + OUTPUTS = ("vn",) + + @pytest.fixture + def input_data(self, mesh): + area_edge = random_field(mesh, EdgeDim) + kh_smag_e = random_field(mesh, EdgeDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) + diff_multfac_vn = random_field(mesh, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + return dict( + area_edge=area_edge, + kh_smag_e=kh_smag_e, + z_nabla2_e=z_nabla2_e, + z_nabla4_e2=z_nabla4_e2, + diff_multfac_vn=diff_multfac_vn, + vn=vn, + ) + + @staticmethod + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py new file mode 100644 index 0000000000..4b4a9ceccf --- /dev/null +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -0,0 +1,69 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( + apply_nabla2_and_nabla4_to_vn, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestApplyNabla2AndNabla4ToVn(StencilTest): + PROGRAM = apply_nabla2_and_nabla4_to_vn + OUTPUTS = ("vn",) + + @pytest.fixture + def input_data(self, mesh): + area_edge = random_field(mesh, EdgeDim) + kh_smag_e = random_field(mesh, EdgeDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) + diff_multfac_vn = random_field(mesh, KDim) + nudgecoeff_e = random_field(mesh, EdgeDim) + vn = random_field(mesh, EdgeDim, KDim) + nudgezone_diff = 9.0 + + return dict( + area_edge=area_edge, + kh_smag_e=kh_smag_e, + z_nabla2_e=z_nabla2_e, + z_nabla4_e2=z_nabla4_e2, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=nudgecoeff_e, + vn=vn, + nudgezone_diff=nudgezone_diff, + ) + + @staticmethod + def reference( + mesh, + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) + vn = vn + area_edge * ( + np.maximum(nudgezone_diff * nudgecoeff_e, kh_smag_e) * z_nabla2_e + - diff_multfac_vn * z_nabla4_e2 * area_edge + ) + return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py new file mode 100644 index 0000000000..7928824322 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -0,0 +1,47 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( + apply_nabla2_to_vn_in_lateral_boundary, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestApplyNabla2ToVnInLateralBoundary(StencilTest): + PROGRAM = apply_nabla2_to_vn_in_lateral_boundary + OUTPUTS = ("vn",) + + @pytest.fixture + def input_data(self, mesh): + fac_bdydiff_v = 5.0 + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + area_edge = random_field(mesh, EdgeDim) + vn = random_field(mesh, EdgeDim, KDim) + return dict( + fac_bdydiff_v=fac_bdydiff_v, + z_nabla2_e=z_nabla2_e, + area_edge=area_edge, + vn=vn, + ) + + @staticmethod + def reference( + mesh, z_nabla2_e: np.array, area_edge: np.array, vn: np.array, fac_bdydiff_v + ) -> np.array: + area_edge = np.expand_dims(area_edge, axis=-1) + vn = vn + (z_nabla2_e * area_edge * fac_bdydiff_v) + return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py new file mode 100644 index 0000000000..c3f4725214 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py @@ -0,0 +1,54 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoApplyNabla2ToW(StencilTest): + PROGRAM = apply_nabla2_to_w + OUTPUTS = ("w",) + + @staticmethod + def reference( + mesh, + area: np.array, + z_nabla2_c: np.array, + geofac_n2s: np.array, + w: np.array, + diff_multfac_w, + ) -> np.array: + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + area = np.expand_dims(area, axis=-1) + w = w - diff_multfac_w * area * area * np.sum(z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1) + return dict(w=w) + + @pytest.fixture + def input_data(self, mesh): + area = random_field(mesh, CellDim) + z_nabla2_c = random_field(mesh, CellDim, KDim) + geofac_n2s = random_field(mesh, CellDim, C2E2CODim) + w = random_field(mesh, CellDim, KDim) + diff_multfac_w = 5.0 + + return dict( + area=area, + z_nabla2_c=z_nabla2_c, + geofac_n2s=geofac_n2s, + w=w, + diff_multfac_w=diff_multfac_w, + ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py new file mode 100644 index 0000000000..98193e0585 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -0,0 +1,52 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( + apply_nabla2_to_w_in_upper_damping_layer, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestApplyNabla2ToWInUpperDampingLayer(StencilTest): + PROGRAM = apply_nabla2_to_w_in_upper_damping_layer + OUTPUTS = ("w",) + + @pytest.fixture + def input_data(self, mesh): + w = random_field(mesh, CellDim, KDim) + diff_multfac_n2w = random_field(mesh, KDim) + cell_area = random_field(mesh, CellDim) + z_nabla2_c = random_field(mesh, CellDim, KDim) + + return dict( + w=w, + diff_multfac_n2w=diff_multfac_n2w, + cell_area=cell_area, + z_nabla2_c=z_nabla2_c, + ) + + @staticmethod + def reference( + mesh, + w: np.array, + diff_multfac_n2w: np.array, + cell_area: np.array, + z_nabla2_c: np.array, + ) -> np.array: + cell_area = np.expand_dims(cell_area, axis=-1) + w = w + diff_multfac_n2w * cell_area * z_nabla2_c + return dict(w=w) diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py new file mode 100644 index 0000000000..c60af2834b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( + calculate_diagnostics_for_turbulence, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestCalculateDiagnosticsForTurbulence(StencilTest): + PROGRAM = calculate_diagnostics_for_turbulence + OUTPUTS = ("div_ic", "hdef_ic") + + @staticmethod + def reference( + mesh, wgtfac_c: np.array, div: np.array, kh_c: np.array, div_ic, hdef_ic + ) -> tuple[np.array, np.array]: + kc_offset_1 = np.roll(kh_c, shift=1, axis=1) + div_offset_1 = np.roll(div, shift=1, axis=1) + div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] + return dict(div_ic=div_ic, hdef_ic=hdef_ic) + + @pytest.fixture + def input_data(self, mesh): + wgtfac_c = random_field(mesh, CellDim, KDim) + div = random_field(mesh, CellDim, KDim) + kh_c = random_field(mesh, CellDim, KDim) + div_ic = zero_field(mesh, CellDim, KDim) + hdef_ic = zero_field(mesh, CellDim, KDim) + return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) diff --git a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py new file mode 100644 index 0000000000..58e285bc0c --- /dev/null +++ b/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -0,0 +1,57 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( + calculate_horizontal_gradients_for_turbulence, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestCalculateHorizontalGradientsForTurbulence(StencilTest): + PROGRAM = calculate_horizontal_gradients_for_turbulence + OUTPUTS = ("dwdx", "dwdy") + + @staticmethod + def reference( + mesh, w: np.array, geofac_grg_x: np.array, geofac_grg_y: np.array, **kwargs + ) -> tuple[np.array]: + geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) + dwdx = np.sum(geofac_grg_x * w[mesh.c2e2cO], axis=1) + + geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) + dwdy = np.sum(geofac_grg_y * w[mesh.c2e2cO], axis=1) + return dict(dwdx=dwdx, dwdy=dwdy) + + @pytest.fixture + def input_data(self, mesh): + w = random_field(mesh, CellDim, KDim) + geofac_grg_x = random_field(mesh, CellDim, C2E2CODim) + geofac_grg_y = random_field(mesh, CellDim, C2E2CODim) + dwdx = zero_field(mesh, CellDim, KDim) + dwdy = zero_field(mesh, CellDim, KDim) + + return dict( + w=w, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + dwdx=dwdx, + dwdy=dwdy, + ) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 94% rename from atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 6861e276f0..06a222e873 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,13 +14,22 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( +from icon4py.model.atmosphere.dycore.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) -from icon4py.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim - -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import ( + E2C2VDim, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def calculate_nabla2_and_smag_coefficients_for_vn_numpy( diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py new file mode 100644 index 0000000000..7fa54ee8ca --- /dev/null +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import ( + calculate_nabla2_for_w, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestCalculateNabla2ForW(StencilTest): + PROGRAM = calculate_nabla2_for_w + OUTPUTS = ("z_nabla2_c",) + + @staticmethod + def reference(mesh, w: np.array, geofac_n2s: np.array, **kwargs) -> np.array: + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + z_nabla2_c = np.sum(w[mesh.c2e2cO] * geofac_n2s, axis=1) + return dict(z_nabla2_c=z_nabla2_c) + + @pytest.fixture + def input_data(self, mesh): + w = random_field(mesh, CellDim, KDim) + geofac_n2s = random_field(mesh, CellDim, C2E2CODim) + z_nabla2_c = zero_field(mesh, CellDim, KDim) + + return dict( + w=w, + geofac_n2s=geofac_n2s, + z_nabla2_c=z_nabla2_c, + ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py new file mode 100644 index 0000000000..4b575c47d5 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py @@ -0,0 +1,56 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( + calculate_nabla2_for_z, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestCalculateNabla2ForZ(StencilTest): + PROGRAM = calculate_nabla2_for_z + OUTPUTS = ("z_nabla2_e",) + + @staticmethod + def reference( + mesh, + kh_smag_e: np.array, + inv_dual_edge_length: np.array, + theta_v: np.array, + **kwargs, + ) -> np.array: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + + theta_v_e2c = theta_v[mesh.e2c] + theta_v_weighted = theta_v_e2c[:, 1] - theta_v_e2c[:, 0] + + z_nabla2_e = kh_smag_e * inv_dual_edge_length * theta_v_weighted + return dict(z_nabla2_e=z_nabla2_e) + + @pytest.fixture + def input_data(self, mesh): + kh_smag_e = random_field(mesh, EdgeDim, KDim) + inv_dual_edge_length = random_field(mesh, EdgeDim) + theta_v = random_field(mesh, CellDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + + return dict( + kh_smag_e=kh_smag_e, + inv_dual_edge_length=inv_dual_edge_length, + theta_v=theta_v, + z_nabla2_e=z_nabla2_e, + ) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py similarity index 77% rename from atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py index 64327ec863..0d6cb63d3a 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py @@ -14,13 +14,16 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( calculate_nabla2_of_theta, ) -from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def calculate_nabla2_of_theta_numpy( @@ -40,9 +43,7 @@ def test_calculate_nabla2_of_theta(): out = zero_field(mesh, CellDim, KDim) - ref = calculate_nabla2_of_theta_numpy( - mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div) - ) + ref = calculate_nabla2_of_theta_numpy(mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div)) calculate_nabla2_of_theta( z_nabla2_e, geofac_div_new, diff --git a/atm_dyn_iconam/tests/test_calculate_nabla4.py b/model/atmosphere/dycore/tests/test_calculate_nabla4.py similarity index 90% rename from atm_dyn_iconam/tests/test_calculate_nabla4.py rename to model/atmosphere/dycore/tests/test_calculate_nabla4.py index e1c61f4995..c2a13b8ef9 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla4.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla4.py @@ -14,11 +14,20 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_nabla4 import calculate_nabla4 -from icon4py.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim - -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.atmosphere.dycore.calculate_nabla4 import calculate_nabla4 +from icon4py.model.common.dimension import ( + E2C2VDim, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def calculate_nabla4_numpy( diff --git a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py new file mode 100644 index 0000000000..50e92d4f58 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -0,0 +1,45 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( + enhance_diffusion_coefficient_for_grid_point_cold_pools, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestEnhanceDiffusionCoefficientForGridPointColdPools(StencilTest): + PROGRAM = enhance_diffusion_coefficient_for_grid_point_cold_pools + OUTPUTS = ("kh_smag_e",) + + @staticmethod + def reference( + mesh, + kh_smag_e: np.array, + enh_diffu_3d: np.array, + ) -> np.array: + kh_smag_e = np.maximum(kh_smag_e, np.max(enh_diffu_3d[mesh.e2c], axis=1)) + return dict(kh_smag_e=kh_smag_e) + + @pytest.fixture + def input_data(self, mesh): + kh_smag_e = random_field(mesh, EdgeDim, KDim) + enh_diffu_3d = random_field(mesh, CellDim, KDim) + + return dict( + kh_smag_e=kh_smag_e, + enh_diffu_3d=enh_diffu_3d, + ) diff --git a/atm_dyn_iconam/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 83% rename from atm_dyn_iconam/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 8f8bc3ba56..a38c366a5c 100644 --- a/atm_dyn_iconam/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -15,13 +15,16 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.mo_advection_traj_btraj_compute_o1_dsl import ( +from icon4py.model.atmosphere.dycore.mo_advection_traj_btraj_compute_o1_dsl import ( mo_advection_traj_btraj_compute_o1_dsl, ) -from icon4py.common.dimension import E2CDim, ECDim, EdgeDim, KDim - -from .test_utils.helpers import as_1D_sparse_field, constant_field, random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + constant_field, + random_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_advection_traj_btraj_compute_o1_dsl_numpy( @@ -51,28 +54,22 @@ def mo_advection_traj_btraj_compute_o1_dsl_numpy( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 diff --git a/model/atmosphere/dycore/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py new file mode 100644 index 0000000000..5323efba3a --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( + mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, +) +from icon4py.model.common.dimension import CellDim, KDim, V2CDim, VertexDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoIconInterpolationScalarCells2vertsScalarRiDsl(StencilTest): + PROGRAM = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl + OUTPUTS = ("p_vert_out",) + + @staticmethod + def reference(mesh, p_cell_in: np.array, c_intp: np.array, **kwargs) -> np.array: + c_intp = np.expand_dims(c_intp, axis=-1) + p_vert_out = np.sum(p_cell_in[mesh.v2c] * c_intp, axis=1) + return dict(p_vert_out=p_vert_out) + + @pytest.fixture + def input_data(self, mesh): + p_cell_in = random_field(mesh, CellDim, KDim) + c_intp = random_field(mesh, VertexDim, V2CDim) + p_vert_out = zero_field(mesh, VertexDim, KDim) + + return dict( + p_cell_in=p_cell_in, + c_intp=c_intp, + p_vert_out=p_vert_out, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py new file mode 100644 index 0000000000..694f94df0d --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -0,0 +1,58 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( + mo_intp_rbf_rbf_vec_interpol_vertex, +) +from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): + PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex + OUTPUTS = ("p_u_out", "p_v_out") + + @staticmethod + def reference( + mesh, p_e_in: np.array, ptr_coeff_1: np.array, ptr_coeff_2: np.array, **kwargs + ) -> tuple[np.array]: + ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) + p_u_out = np.sum(p_e_in[mesh.v2e] * ptr_coeff_1, axis=1) + + ptr_coeff_2 = np.expand_dims(ptr_coeff_2, axis=-1) + p_v_out = np.sum(p_e_in[mesh.v2e] * ptr_coeff_2, axis=1) + + return dict(p_v_out=p_v_out, p_u_out=p_u_out) + + @pytest.fixture + def input_data(self, mesh): + p_e_in = random_field(mesh, EdgeDim, KDim) + ptr_coeff_1 = random_field(mesh, VertexDim, V2EDim) + ptr_coeff_2 = random_field(mesh, VertexDim, V2EDim) + p_v_out = zero_field(mesh, VertexDim, KDim) + p_u_out = zero_field(mesh, VertexDim, KDim) + + return dict( + p_e_in=p_e_in, + ptr_coeff_1=ptr_coeff_1, + ptr_coeff_2=ptr_coeff_2, + p_v_out=p_v_out, + p_u_out=p_u_out, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py new file mode 100644 index 0000000000..ae29a9bbf7 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_math_divrot_rot_vertex_ri_dsl import ( + mo_math_divrot_rot_vertex_ri_dsl, +) +from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoMathDivrotRotVertexRiDsl(StencilTest): + PROGRAM = mo_math_divrot_rot_vertex_ri_dsl + OUTPUTS = ("rot_vec",) + + @staticmethod + def reference(mesh, vec_e: np.array, geofac_rot: np.array, **kwargs) -> np.array: + geofac_rot = np.expand_dims(geofac_rot, axis=-1) + rot_vec = np.sum(vec_e[mesh.v2e] * geofac_rot, axis=1) + return dict(rot_vec=rot_vec) + + @pytest.fixture + def input_data(self, mesh): + vec_e = random_field(mesh, EdgeDim, KDim) + geofac_rot = random_field(mesh, VertexDim, V2EDim) + rot_vec = zero_field(mesh, VertexDim, KDim) + + return dict( + vec_e=vec_e, + geofac_rot=geofac_rot, + rot_vec=rot_vec, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py b/model/atmosphere/dycore/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py new file mode 100644 index 0000000000..e13c3d923b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py @@ -0,0 +1,77 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_math_gradients_grad_green_gauss_cell_dsl import ( + mo_math_gradients_grad_green_gauss_cell_dsl, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoMathGradientsGradGreenGaussCellDsl(StencilTest): + PROGRAM = mo_math_gradients_grad_green_gauss_cell_dsl + OUTPUTS = ("p_grad_1_u", "p_grad_1_v", "p_grad_2_u", "p_grad_2_v") + + @staticmethod + def reference( + mesh, + p_ccpr1: np.array, + p_ccpr2: np.array, + geofac_grg_x: np.array, + geofac_grg_y: np.array, + **kwargs, + ) -> tuple[np.array]: + geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) + p_grad_1_u = np.sum(geofac_grg_x * p_ccpr1[mesh.c2e2cO], axis=1) + + geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) + p_grad_1_v = np.sum(geofac_grg_y * p_ccpr1[mesh.c2e2cO], axis=1) + + p_grad_2_u = np.sum(geofac_grg_x * p_ccpr2[mesh.c2e2cO], axis=1) + + p_grad_2_v = np.sum(geofac_grg_y * p_ccpr2[mesh.c2e2cO], axis=1) + return dict( + p_grad_1_u=p_grad_1_u, + p_grad_1_v=p_grad_1_v, + p_grad_2_u=p_grad_2_u, + p_grad_2_v=p_grad_2_v, + ) + + @pytest.fixture + def input_data(self, mesh): + p_ccpr1 = random_field(mesh, CellDim, KDim) + p_ccpr2 = random_field(mesh, CellDim, KDim) + geofac_grg_x = random_field(mesh, CellDim, C2E2CODim) + geofac_grg_y = random_field(mesh, CellDim, C2E2CODim) + p_grad_1_u = zero_field(mesh, CellDim, KDim) + p_grad_1_v = zero_field(mesh, CellDim, KDim) + p_grad_2_u = zero_field(mesh, CellDim, KDim) + p_grad_2_v = zero_field(mesh, CellDim, KDim) + + return dict( + p_grad_1_u=p_grad_1_u, + p_grad_1_v=p_grad_1_v, + p_grad_2_u=p_grad_2_u, + p_grad_2_v=p_grad_2_v, + p_ccpr1=p_ccpr1, + p_ccpr2=p_ccpr2, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_4th_order_divdamp.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_4th_order_divdamp.py new file mode 100644 index 0000000000..d0c5ef81eb --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_4th_order_divdamp.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_4th_order_divdamp import ( + mo_solve_nonhydro_4th_order_divdamp, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydro4thOrderDivdamp(StencilTest): + PROGRAM = mo_solve_nonhydro_4th_order_divdamp + OUTPUTS = ("vn",) + + @staticmethod + def reference( + mesh, + scal_divdamp: np.array, + z_graddiv2_vn: np.array, + vn: np.array, + ) -> np.array: + scal_divdamp = np.expand_dims(scal_divdamp, axis=0) + vn = vn + (scal_divdamp * z_graddiv2_vn) + return dict(vn=vn) + + @pytest.fixture + def input_data(self, mesh): + scal_divdamp = random_field(mesh, KDim) + z_graddiv2_vn = random_field(mesh, EdgeDim, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + return dict( + scal_divdamp=scal_divdamp, + z_graddiv2_vn=z_graddiv2_vn, + vn=vn, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_01.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_01.py new file mode 100644 index 0000000000..ef7adaf0e8 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_01.py @@ -0,0 +1,43 @@ +# 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 + +import numpy as np +import pytest as pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_01 import ( + mo_solve_nonhydro_stencil_01, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil01(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_01 + OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") + + @staticmethod + def reference( + mesh, + z_rth_pr_1: np.array, + z_rth_pr_2: np.array, + ) -> tuple[np.array]: + z_rth_pr_1 = np.zeros_like(z_rth_pr_1) + z_rth_pr_2 = np.zeros_like(z_rth_pr_2) + return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) + + @pytest.fixture + def input_data(self, mesh): + z_rth_pr_1 = random_field(mesh, CellDim, KDim) + z_rth_pr_2 = random_field(mesh, CellDim, KDim) + + return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py new file mode 100644 index 0000000000..4238b31fd3 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py @@ -0,0 +1,59 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_02 import ( + mo_solve_nonhydro_stencil_02, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil02(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_02 + OUTPUTS = ("z_exner_ex_pr", "exner_pr") + + @staticmethod + def reference( + mesh, + exner: np.array, + exner_ref_mc: np.array, + exner_pr: np.array, + exner_exfac: np.array, + **kwargs, + ) -> dict: + z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + exner_pr = exner - exner_ref_mc + return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) + + @pytest.fixture + def input_data(self, mesh): + exner = random_field(mesh, CellDim, KDim) + exner_ref_mc = random_field(mesh, CellDim, KDim) + exner_pr = zero_field(mesh, CellDim, KDim) + exner_exfac = random_field(mesh, CellDim, KDim) + z_exner_ex_pr = zero_field(mesh, CellDim, KDim) + + return dict( + exner_exfac=exner_exfac, + exner=exner, + exner_ref_mc=exner_ref_mc, + exner_pr=exner_pr, + z_exner_ex_pr=z_exner_ex_pr, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_03.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_03.py new file mode 100644 index 0000000000..e84f0e13d4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_03.py @@ -0,0 +1,39 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_03 import ( + mo_solve_nonhydro_stencil_03, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil03(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_03 + OUTPUTS = ("z_exner_ex_pr",) + + @staticmethod + def reference(mesh, z_exner_ex_pr: np.array, **kwargs) -> dict: + z_exner_ex_pr = np.zeros_like(z_exner_ex_pr) + return dict(z_exner_ex_pr=z_exner_ex_pr) + + @pytest.fixture + def input_data(self, mesh): + z_exner_ex_pr = random_field(mesh, CellDim, KDim) + + return dict( + z_exner_ex_pr=z_exner_ex_pr, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py new file mode 100644 index 0000000000..29cb23efed --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py @@ -0,0 +1,56 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_04 import ( + mo_solve_nonhydro_stencil_04, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil04(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_04 + OUTPUTS = ("z_exner_ic",) + + @staticmethod + def reference( + mesh, + z_exner_ex_pr: np.array, + wgtfacq_c: np.array, + z_exner_ic: np.array, + ) -> np.array: + z_exner_ic[:, 3:] = ( + np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) + + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) + + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) + )[:, 3:] + return {"z_exner_ic": z_exner_ic} + + @pytest.fixture + def input_data(self, mesh): + z_exner_ex_pr = random_field(mesh, CellDim, KDim) + wgtfacq_c = random_field(mesh, CellDim, KDim) + z_exner_ic = zero_field(mesh, CellDim, KDim) + + return dict( + z_exner_ex_pr=z_exner_ex_pr, + wgtfacq_c=wgtfacq_c, + z_exner_ic=z_exner_ic, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py similarity index 83% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_05.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py index 51a52f2b10..857f30cee9 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_05 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_05 import ( mo_solve_nonhydro_stencil_05, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_05_numpy( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py new file mode 100644 index 0000000000..9a86409fa4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py @@ -0,0 +1,47 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_06 import ( + mo_solve_nonhydro_stencil_06, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil06(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_06 + OUTPUTS = ("z_dexner_dz_c_1",) + + @staticmethod + def reference(mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs) -> np.array: + z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full + return dict(z_dexner_dz_c_1=z_dexner_dz_c_1) + + @pytest.fixture + def input_data(self, mesh): + z_exner_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + z_dexner_dz_c_1 = zero_field(mesh, CellDim, KDim) + + return dict( + z_exner_ic=z_exner_ic, + inv_ddqz_z_full=inv_ddqz_z_full, + z_dexner_dz_c_1=z_dexner_dz_c_1, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_07.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_07.py new file mode 100644 index 0000000000..9b3997ce3f --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_07.py @@ -0,0 +1,61 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_07 import ( + mo_solve_nonhydro_stencil_07, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil07(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_07 + OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") + + @staticmethod + def reference( + mesh, + rho: np.array, + rho_ref_mc: np.array, + theta_v: np.array, + theta_ref_mc: np.array, + **kwargs, + ) -> tuple[np.array]: + z_rth_pr_1 = rho - rho_ref_mc + z_rth_pr_2 = theta_v - theta_ref_mc + return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) + + @pytest.fixture + def input_data(self, mesh): + rho = random_field(mesh, CellDim, KDim) + rho_ref_mc = random_field(mesh, CellDim, KDim) + theta_v = random_field(mesh, CellDim, KDim) + theta_ref_mc = random_field(mesh, CellDim, KDim) + z_rth_pr_1 = zero_field(mesh, CellDim, KDim) + z_rth_pr_2 = zero_field(mesh, CellDim, KDim) + + return dict( + rho=rho, + rho_ref_mc=rho_ref_mc, + theta_v=theta_v, + theta_ref_mc=theta_ref_mc, + z_rth_pr_1=z_rth_pr_1, + z_rth_pr_2=z_rth_pr_2, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py similarity index 88% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_08.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py index 430c52b785..2154cb9062 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_08 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_08 import ( mo_solve_nonhydro_stencil_08, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_08_numpy( diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py similarity index 91% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_09.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py index 1312a59464..6d9496c0c9 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_09 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_09 import ( mo_solve_nonhydro_stencil_09, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_09_numpy( diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py similarity index 95% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_10.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py index e7c1c5181a..ad9f5ff2d4 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py @@ -13,13 +13,11 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_10 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_10 import ( mo_solve_nonhydro_stencil_10, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import SimpleMesh, random_field def mo_solve_nonhydro_stencil_10_numpy( @@ -39,7 +37,6 @@ def mo_solve_nonhydro_stencil_10_numpy( wgt_nnow_rth, wgt_nnew_rth, ) -> tuple[np.array, np.array, np.array, np.array]: - vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) rho_now_offset = np.roll(rho_now, shift=1, axis=1) rho_var_offset = np.roll(rho_var, shift=1, axis=1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_lower.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_lower.py new file mode 100644 index 0000000000..5840024652 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_lower.py @@ -0,0 +1,38 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_11_lower import ( + mo_solve_nonhydro_stencil_11_lower, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil11Lower(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_11_lower + OUTPUTS = ("z_theta_v_pr_ic",) + + @staticmethod + def reference(mesh, **kwargs) -> np.array: + z_theta_v_pr_ic = 0 + return dict(z_theta_v_pr_ic=z_theta_v_pr_ic) + + @pytest.fixture + def input_data(self, mesh): + z_theta_v_pr_ic = random_field(mesh, CellDim, KDim) + return dict( + z_theta_v_pr_ic=z_theta_v_pr_ic, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_upper.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py similarity index 87% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_upper.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py index a050e0bf3b..aafd2d688e 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_11_upper.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_11_upper import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_11_upper import ( mo_solve_nonhydro_stencil_11_upper, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_11_upper_numpy( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_12.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_12.py new file mode 100644 index 0000000000..175bc287fb --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_12.py @@ -0,0 +1,63 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_12 import ( + mo_solve_nonhydro_stencil_12, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil12(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_12 + OUTPUTS = ("z_dexner_dz_c_2",) + + @staticmethod + def reference( + mesh, + z_theta_v_pr_ic: np.array, + d2dexdz2_fac1_mc: np.array, + d2dexdz2_fac2_mc: np.array, + z_rth_pr_2: np.array, + **kwargs, + ) -> np.array: + z_theta_v_pr_ic_offset_1 = z_theta_v_pr_ic[:, 1:] + z_dexner_dz_c_2 = -0.5 * ( + (z_theta_v_pr_ic[:, :-1] - z_theta_v_pr_ic_offset_1) * d2dexdz2_fac1_mc + + z_rth_pr_2 * d2dexdz2_fac2_mc + ) + return dict(z_dexner_dz_c_2=z_dexner_dz_c_2) + + @pytest.fixture + def input_data(self, mesh): + z_theta_v_pr_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + d2dexdz2_fac1_mc = random_field(mesh, CellDim, KDim) + z_rth_pr_2 = random_field(mesh, CellDim, KDim) + d2dexdz2_fac2_mc = random_field(mesh, CellDim, KDim) + + z_dexner_dz_c_2 = zero_field(mesh, CellDim, KDim) + + return dict( + z_theta_v_pr_ic=z_theta_v_pr_ic, + d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, + d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + z_rth_pr_2=z_rth_pr_2, + z_dexner_dz_c_2=z_dexner_dz_c_2, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_13.py new file mode 100644 index 0000000000..bf21c4f5f4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_13.py @@ -0,0 +1,61 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_13 import ( + mo_solve_nonhydro_stencil_13, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil13(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_13 + OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") + + @staticmethod + def reference( + mesh, + rho: np.array, + rho_ref_mc: np.array, + theta_v: np.array, + theta_ref_mc: np.array, + **kwargs, + ) -> dict: + z_rth_pr_1 = rho - rho_ref_mc + z_rth_pr_2 = theta_v - theta_ref_mc + return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) + + @pytest.fixture + def input_data(self, mesh): + rho = random_field(mesh, CellDim, KDim) + rho_ref_mc = random_field(mesh, CellDim, KDim) + theta_v = random_field(mesh, CellDim, KDim) + theta_ref_mc = random_field(mesh, CellDim, KDim) + z_rth_pr_1 = zero_field(mesh, CellDim, KDim) + z_rth_pr_2 = zero_field(mesh, CellDim, KDim) + + return dict( + rho=rho, + rho_ref_mc=rho_ref_mc, + theta_v=theta_v, + theta_ref_mc=theta_ref_mc, + z_rth_pr_1=z_rth_pr_1, + z_rth_pr_2=z_rth_pr_2, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_14.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_14.py new file mode 100644 index 0000000000..a6a48b1208 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_14.py @@ -0,0 +1,42 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_14 import ( + mo_solve_nonhydro_stencil_14, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil14(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_14 + OUTPUTS = ("z_rho_e", "z_theta_v_e") + + @staticmethod + def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> dict: + z_rho_e = np.zeros_like(z_rho_e) + z_theta_v_e = np.zeros_like(z_theta_v_e) + return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) + + @pytest.fixture + def input_data(self, mesh): + z_rho_e = zero_field(mesh, EdgeDim, KDim) + z_theta_v_e = zero_field(mesh, EdgeDim, KDim) + + return dict( + z_rho_e=z_rho_e, + z_theta_v_e=z_theta_v_e, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py new file mode 100644 index 0000000000..bb92f1fbef --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py @@ -0,0 +1,42 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_15 import ( + mo_solve_nonhydro_stencil_15, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil15(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_15 + OUTPUTS = ("z_rho_e", "z_theta_v_e") + + @staticmethod + def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> tuple[np.array]: + z_rho_e = np.zeros_like(z_rho_e) + z_theta_v_e = np.zeros_like(z_theta_v_e) + return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) + + @pytest.fixture + def input_data(self, mesh): + z_rho_e = random_field(mesh, EdgeDim, KDim) + z_theta_v_e = random_field(mesh, EdgeDim, KDim) + + return dict( + z_rho_e=z_rho_e, + z_theta_v_e=z_theta_v_e, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py similarity index 87% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index e4ab4303cc..8282ea91e1 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -14,13 +14,15 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1 import ( mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1, ) -from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim - -from .test_utils.helpers import as_1D_sparse_field, random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def compute_btraj_numpy( @@ -43,28 +45,22 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_distv_bary_1, p_distv_bary_2 @@ -84,7 +80,6 @@ def sten_16_numpy( z_rth_pr_1: np.array, z_rth_pr_2: np.array, ) -> np.array: - z_rth_pr_1_e2c = z_rth_pr_1[e2c] z_rth_pr_2_e2c = z_rth_pr_2[e2c] z_grad_rth_1_e2c = z_grad_rth_1[e2c] @@ -199,10 +194,7 @@ def test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1(): z_rho_e = random_field(mesh, EdgeDim, KDim) z_theta_v_e = random_field(mesh, EdgeDim, KDim) - ( - z_rho_e_ref, - z_theta_v_e_ref, - ) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( + (z_rho_e_ref, z_theta_v_e_ref,) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( mesh.e2c, np.asarray(p_vn), np.asarray(p_vt), diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_17.py new file mode 100644 index 0000000000..a770c892b4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_17.py @@ -0,0 +1,64 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_17 import ( + mo_solve_nonhydro_stencil_17, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil17(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_17 + OUTPUTS = ("z_graddiv_vn",) + + @staticmethod + def reference( + mesh, + hmask_dd3d: np.array, + scalfac_dd3d: np.array, + inv_dual_edge_length: np.array, + z_dwdz_dd: np.array, + z_graddiv_vn: np.array, + **kwargs, + ) -> dict: + scalfac_dd3d = np.expand_dims(scalfac_dd3d, axis=0) + hmask_dd3d = np.expand_dims(hmask_dd3d, axis=-1) + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + + z_dwdz_dd_e2c = z_dwdz_dd[mesh.e2c] + z_dwdz_dd_weighted = z_dwdz_dd_e2c[:, 1] - z_dwdz_dd_e2c[:, 0] + + z_graddiv_vn = z_graddiv_vn + ( + hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * z_dwdz_dd_weighted + ) + return dict(z_graddiv_vn=z_graddiv_vn) + + @pytest.fixture + def input_data(self, mesh): + hmask_dd3d = random_field(mesh, EdgeDim) + scalfac_dd3d = random_field(mesh, KDim) + inv_dual_edge_length = random_field(mesh, EdgeDim) + z_dwdz_dd = random_field(mesh, CellDim, KDim) + z_graddiv_vn = random_field(mesh, EdgeDim, KDim) + + return dict( + hmask_dd3d=hmask_dd3d, + scalfac_dd3d=scalfac_dd3d, + inv_dual_edge_length=inv_dual_edge_length, + z_dwdz_dd=z_dwdz_dd, + z_graddiv_vn=z_graddiv_vn, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_18.py new file mode 100644 index 0000000000..d04c4e1b24 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_18.py @@ -0,0 +1,50 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_18 import ( + mo_solve_nonhydro_stencil_18, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil18(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_18 + OUTPUTS = ("z_gradh_exner",) + + @staticmethod + def reference( + mesh, inv_dual_edge_length: np.array, z_exner_ex_pr: np.array, **kwargs + ) -> np.array: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + + z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] + z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] + + z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted + return dict(z_gradh_exner=z_gradh_exner) + + @pytest.fixture + def input_data(self, mesh): + inv_dual_edge_length = random_field(mesh, EdgeDim) + z_exner_ex_pr = random_field(mesh, CellDim, KDim) + z_gradh_exner = random_field(mesh, EdgeDim, KDim) + + return dict( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=z_exner_ex_pr, + z_gradh_exner=z_gradh_exner, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py new file mode 100644 index 0000000000..5a782f7c5f --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py @@ -0,0 +1,65 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_19 import ( + mo_solve_nonhydro_stencil_19, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil19(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_19 + OUTPUTS = ("z_gradh_exner",) + + @staticmethod + def reference( + mesh, + inv_dual_edge_length: np.array, + z_exner_ex_pr: np.array, + ddxn_z_full: np.array, + c_lin_e: np.array, + z_dexner_dz_c_1: np.array, + **kwargs, + ) -> dict: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + + z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] + z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] + + z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( + c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1 + ) + return dict(z_gradh_exner=z_gradh_exner) + + @pytest.fixture + def input_data(self, mesh): + inv_dual_edge_length = random_field(mesh, EdgeDim) + z_exner_ex_pr = random_field(mesh, CellDim, KDim) + ddxn_z_full = random_field(mesh, EdgeDim, KDim) + c_lin_e = random_field(mesh, EdgeDim, E2CDim) + z_dexner_dz_c_1 = random_field(mesh, CellDim, KDim) + z_gradh_exner = random_field(mesh, EdgeDim, KDim) + + return dict( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=z_exner_ex_pr, + ddxn_z_full=ddxn_z_full, + c_lin_e=c_lin_e, + z_dexner_dz_c_1=z_dexner_dz_c_1, + z_gradh_exner=z_gradh_exner, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py similarity index 85% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py index 9ec1320cd0..349355cd72 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py @@ -15,13 +15,16 @@ 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 ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, ) -from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim - -from .test_utils.helpers import flatten_first_two_dims, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + flatten_first_two_dims, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_20_numpy( @@ -48,12 +51,8 @@ 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, :] * ( @@ -101,11 +100,6 @@ def test_mo_solve_nonhydro_stencil_20(): np.asarray(z_dexner_dz_c_2), ) - hstart = 0 - hend = mesh.n_edges - kstart = 0 - kend = mesh.k_level - mo_solve_nonhydro_stencil_20( inv_dual_edge_length, z_exner_ex_pr, @@ -114,10 +108,6 @@ def test_mo_solve_nonhydro_stencil_20(): z_dexner_dz_c_1, z_dexner_dz_c_2, z_gradh_exner, - hstart, - hend, - kstart, - kend, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py similarity index 89% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py index cf2d4ee3d9..e0f31fb5ef 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py @@ -15,13 +15,16 @@ 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 ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_21 import ( mo_solve_nonhydro_stencil_21, ) -from icon4py.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim - -from .test_utils.helpers import flatten_first_two_dims, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + flatten_first_two_dims, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_21_numpy( @@ -58,9 +61,7 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): full_shape, theta_v_ic, e2c, ikoffset ) - inv_ddqz_z_full_at_kidx, _ = _apply_index_field( - full_shape, inv_ddqz_z_full, e2c, ikoffset - ) + inv_ddqz_z_full_at_kidx, _ = _apply_index_field(full_shape, inv_ddqz_z_full, e2c, ikoffset) z_theta1 = ( theta_v_at_kidx[:, 0, :] @@ -123,11 +124,6 @@ def test_mo_solve_nonhydro_stencil_21(): grav_o_cpd, ) - hstart = 0 - hend = mesh.n_edges - kstart = 0 - kend = mesh.k_level - mo_solve_nonhydro_stencil_21( theta_v, ikoffset_new, @@ -137,10 +133,6 @@ def test_mo_solve_nonhydro_stencil_21(): inv_dual_edge_length, grav_o_cpd, z_hydro_corr, - hstart, - hend, - kstart, - kend, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_22.py new file mode 100644 index 0000000000..5273d7d3d8 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_22.py @@ -0,0 +1,59 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_22 import ( + mo_solve_nonhydro_stencil_22, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, +) + + +class TestMoSolveNonhydroStencil22(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_22 + OUTPUTS = ("z_gradh_exner",) + + @staticmethod + def reference( + mesh, + ipeidx_dsl: np.array, + pg_exdist: np.array, + z_hydro_corr: np.array, + z_gradh_exner: np.array, + **kwargs, + ) -> dict: + z_hydro_corr = np.expand_dims(z_hydro_corr, axis=-1) + z_gradh_exner = np.where( + ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner + ) + return dict(z_gradh_exner=z_gradh_exner) + + @pytest.fixture + def input_data(self, mesh): + ipeidx_dsl = random_mask(mesh, EdgeDim, KDim) + pg_exdist = random_field(mesh, EdgeDim, KDim) + z_hydro_corr = random_field(mesh, EdgeDim) + z_gradh_exner = random_field(mesh, EdgeDim, KDim) + + return dict( + ipeidx_dsl=ipeidx_dsl, + pg_exdist=pg_exdist, + z_hydro_corr=z_hydro_corr, + z_gradh_exner=z_gradh_exner, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_23.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_23.py new file mode 100644 index 0000000000..388938caf9 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_23.py @@ -0,0 +1,81 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_23 import ( + mo_solve_nonhydro_stencil_23, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil23(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_23 + OUTPUTS = ("vn_nnew",) + + @staticmethod + def reference( + mesh, + vn_nnow: np.array, + ddt_vn_adv_ntl1: np.array, + ddt_vn_adv_ntl2: np.array, + ddt_vn_phy: np.array, + z_theta_v_e: np.array, + z_gradh_exner: np.array, + dtime, + wgt_nnow_vel, + wgt_nnew_vel, + cpd, + **kwargs, + ) -> np.array: + vn_nnew = vn_nnow + dtime * ( + wgt_nnow_vel * ddt_vn_adv_ntl1 + + wgt_nnew_vel * ddt_vn_adv_ntl2 + + ddt_vn_phy + - cpd * z_theta_v_e * z_gradh_exner + ) + return dict(vn_nnew=vn_nnew) + + @pytest.fixture + def input_data(self, mesh): + vn_nnow = random_field(mesh, EdgeDim, KDim) + ddt_vn_adv_ntl1 = random_field(mesh, EdgeDim, KDim) + ddt_vn_adv_ntl2 = random_field(mesh, EdgeDim, KDim) + ddt_vn_phy = random_field(mesh, EdgeDim, KDim) + z_theta_v_e = random_field(mesh, EdgeDim, KDim) + z_gradh_exner = random_field(mesh, EdgeDim, KDim) + vn_nnew = zero_field(mesh, EdgeDim, KDim) + dtime = 5.0 + wgt_nnow_vel = 8.0 + wgt_nnew_vel = 7.0 + cpd = 2.0 + + return dict( + vn_nnow=vn_nnow, + ddt_vn_adv_ntl1=ddt_vn_adv_ntl1, + ddt_vn_adv_ntl2=ddt_vn_adv_ntl2, + ddt_vn_phy=ddt_vn_phy, + z_theta_v_e=z_theta_v_e, + z_gradh_exner=z_gradh_exner, + vn_nnew=vn_nnew, + dtime=dtime, + wgt_nnow_vel=wgt_nnow_vel, + wgt_nnew_vel=wgt_nnew_vel, + cpd=cpd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_24.py new file mode 100644 index 0000000000..1887cad8a8 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_24.py @@ -0,0 +1,68 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_24 import ( + mo_solve_nonhydro_stencil_24, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil24(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_24 + OUTPUTS = ("vn_nnew",) + + @staticmethod + def reference( + mesh, + vn_nnow: np.array, + ddt_vn_adv_ntl1: np.array, + ddt_vn_phy: np.array, + z_theta_v_e: np.array, + z_gradh_exner: np.array, + dtime: float, + cpd: float, + **kwargs, + ) -> np.array: + vn_nnew = vn_nnow + dtime * ( + ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner + ) + return dict(vn_nnew=vn_nnew) + + @pytest.fixture + def input_data(self, mesh): + dtime, cpd = 10.0, 10.0 + vn_nnow = random_field(mesh, EdgeDim, KDim) + ddt_vn_adv_ntl1 = random_field(mesh, EdgeDim, KDim) + ddt_vn_phy = random_field(mesh, EdgeDim, KDim) + z_theta_v_e = random_field(mesh, EdgeDim, KDim) + z_gradh_exner = random_field(mesh, EdgeDim, KDim) + vn_nnew = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn_nnow=vn_nnow, + ddt_vn_adv_ntl1=ddt_vn_adv_ntl1, + ddt_vn_phy=ddt_vn_phy, + z_theta_v_e=z_theta_v_e, + z_gradh_exner=z_gradh_exner, + vn_nnew=vn_nnew, + dtime=dtime, + cpd=cpd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py new file mode 100644 index 0000000000..18623f00db --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_25 import ( + mo_solve_nonhydro_stencil_25, +) +from icon4py.model.common.dimension import E2C2EODim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil25(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_25 + OUTPUTS = ("z_graddiv2_vn",) + + @staticmethod + def reference(mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs) -> np.array: + geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) + z_graddiv2_vn = np.sum(z_graddiv_vn[mesh.e2c2eO] * geofac_grdiv, axis=1) + return dict(z_graddiv2_vn=z_graddiv2_vn) + + @pytest.fixture + def input_data(self, mesh): + z_graddiv_vn = random_field(mesh, EdgeDim, KDim) + geofac_grdiv = random_field(mesh, EdgeDim, E2C2EODim) + z_graddiv2_vn = zero_field(mesh, EdgeDim, KDim) + + return dict( + geofac_grdiv=geofac_grdiv, + z_graddiv_vn=z_graddiv_vn, + z_graddiv2_vn=z_graddiv2_vn, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py new file mode 100644 index 0000000000..ec2389a8de --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py @@ -0,0 +1,43 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_26 import ( + mo_solve_nonhydro_stencil_26, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil26(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_26 + OUTPUTS = ("vn",) + + @staticmethod + def reference(mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs) -> dict: + vn = vn + (scal_divdamp_o2 * z_graddiv_vn) + return dict(vn=vn) + + @pytest.fixture + def input_data(self, mesh): + z_graddiv_vn = random_field(mesh, EdgeDim, KDim) + vn = random_field(mesh, EdgeDim, KDim) + scal_divdamp_o2 = 5.0 + + return dict( + z_graddiv_vn=z_graddiv_vn, + vn=vn, + scal_divdamp_o2=scal_divdamp_o2, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_27.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_27.py new file mode 100644 index 0000000000..faaf543ed0 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_27.py @@ -0,0 +1,56 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_27 import ( + mo_solve_nonhydro_stencil_27, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil27(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_27 + OUTPUTS = ("vn",) + + @staticmethod + def reference( + mesh, + scal_divdamp: np.array, + bdy_divdamp: np.array, + nudgecoeff_e: np.array, + z_graddiv2_vn: np.array, + vn: np.array, + **kwargs, + ) -> dict: + nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) + vn = vn + (scal_divdamp + bdy_divdamp * nudgecoeff_e) * z_graddiv2_vn + return dict(vn=vn) + + @pytest.fixture + def input_data(self, mesh): + scal_divdamp = random_field(mesh, KDim) + bdy_divdamp = random_field(mesh, KDim) + nudgecoeff_e = random_field(mesh, EdgeDim) + z_graddiv2_vn = random_field(mesh, EdgeDim, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + return dict( + scal_divdamp=scal_divdamp, + bdy_divdamp=bdy_divdamp, + nudgecoeff_e=nudgecoeff_e, + z_graddiv2_vn=z_graddiv2_vn, + vn=vn, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py new file mode 100644 index 0000000000..17234d448c --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py @@ -0,0 +1,43 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_28 import ( + mo_solve_nonhydro_stencil_28, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil28(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_28 + OUTPUTS = ("vn",) + + @staticmethod + def reference(mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs) -> np.array: + vn = vn + (iau_wgt_dyn * vn_incr) + return dict(vn=vn) + + @pytest.fixture + def input_data(self, mesh): + vn_incr = random_field(mesh, EdgeDim, KDim) + vn = random_field(mesh, EdgeDim, KDim) + iau_wgt_dyn = 5.0 + + return dict( + vn_incr=vn_incr, + vn=vn, + iau_wgt_dyn=iau_wgt_dyn, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py new file mode 100644 index 0000000000..da69d23da3 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_29 import ( + mo_solve_nonhydro_stencil_29, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil29(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_29 + OUTPUTS = ("vn_new",) + + @staticmethod + def reference(mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs) -> dict: + vn_new = vn_now + dtime * grf_tend_vn + return dict(vn_new=vn_new) + + @pytest.fixture + def input_data(self, mesh): + grf_tend_vn = random_field(mesh, EdgeDim, KDim) + vn_now = random_field(mesh, EdgeDim, KDim) + vn_new = zero_field(mesh, EdgeDim, KDim) + dtime = 6.0 + + return dict( + grf_tend_vn=grf_tend_vn, + vn_now=vn_now, + vn_new=vn_new, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_30.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_30.py new file mode 100644 index 0000000000..200122bf07 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_30.py @@ -0,0 +1,67 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_30 import ( + mo_solve_nonhydro_stencil_30, +) +from icon4py.model.common.dimension import E2C2EDim, E2C2EODim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil30(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_30 + OUTPUTS = ("z_vn_avg", "z_graddiv_vn", "vt") + + @staticmethod + def reference( + mesh, + e_flx_avg: np.array, + vn: np.array, + geofac_grdiv: np.array, + rbf_vec_coeff_e: np.array, + **kwargs, + ) -> dict: + e_flx_avg = np.expand_dims(e_flx_avg, axis=-1) + z_vn_avg = np.sum(vn[mesh.e2c2eO] * e_flx_avg, axis=1) + geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) + z_graddiv_vn = np.sum(vn[mesh.e2c2eO] * geofac_grdiv, axis=1) + rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) + vt = np.sum(vn[mesh.e2c2e] * rbf_vec_coeff_e, axis=1) + return dict(z_vn_avg=z_vn_avg, z_graddiv_vn=z_graddiv_vn, vt=vt) + + @pytest.fixture + def input_data(self, mesh): + e_flx_avg = random_field(mesh, EdgeDim, E2C2EODim) + geofac_grdiv = random_field(mesh, EdgeDim, E2C2EODim) + rbf_vec_coeff_e = random_field(mesh, EdgeDim, E2C2EDim) + vn = random_field(mesh, EdgeDim, KDim) + z_vn_avg = zero_field(mesh, EdgeDim, KDim) + z_graddiv_vn = zero_field(mesh, EdgeDim, KDim) + vt = zero_field(mesh, EdgeDim, KDim) + + return dict( + e_flx_avg=e_flx_avg, + vn=vn, + geofac_grdiv=geofac_grdiv, + rbf_vec_coeff_e=rbf_vec_coeff_e, + z_vn_avg=z_vn_avg, + z_graddiv_vn=z_graddiv_vn, + vt=vt, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_31.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_31.py new file mode 100644 index 0000000000..0722849c37 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_31.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_31 import ( + mo_solve_nonhydro_stencil_31, +) +from icon4py.model.common.dimension import E2C2EODim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil31(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_31 + OUTPUTS = ("z_vn_avg",) + + @staticmethod + def reference(mesh, e_flx_avg: np.array, vn: np.array, **kwargs) -> np.array: + geofac_grdiv = np.expand_dims(e_flx_avg, axis=-1) + z_vn_avg = np.sum(vn[mesh.e2c2eO] * geofac_grdiv, axis=1) + return dict(z_vn_avg=z_vn_avg) + + @pytest.fixture + def input_data(self, mesh): + e_flx_avg = random_field(mesh, EdgeDim, E2C2EODim) + vn = random_field(mesh, EdgeDim, KDim) + z_vn_avg = zero_field(mesh, EdgeDim, KDim) + + return dict( + e_flx_avg=e_flx_avg, + vn=vn, + z_vn_avg=z_vn_avg, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_32.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_32.py new file mode 100644 index 0000000000..9050261161 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_32.py @@ -0,0 +1,61 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_32 import ( + mo_solve_nonhydro_stencil_32, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil32(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_32 + OUTPUTS = ("mass_fl_e", "z_theta_v_fl_e") + + @staticmethod + def reference( + mesh, + z_rho_e: np.array, + z_vn_avg: np.array, + ddqz_z_full_e: np.array, + z_theta_v_e: np.array, + **kwargs, + ) -> dict: + mass_fl_e = z_rho_e * z_vn_avg * ddqz_z_full_e + z_theta_v_fl_e = mass_fl_e * z_theta_v_e + return dict(mass_fl_e=mass_fl_e, z_theta_v_fl_e=z_theta_v_fl_e) + + @pytest.fixture + def input_data(self, mesh): + z_rho_e = random_field(mesh, EdgeDim, KDim) + z_vn_avg = random_field(mesh, EdgeDim, KDim) + ddqz_z_full_e = random_field(mesh, EdgeDim, KDim) + mass_fl_e = random_field(mesh, EdgeDim, KDim) + z_theta_v_e = random_field(mesh, EdgeDim, KDim) + z_theta_v_fl_e = zero_field(mesh, EdgeDim, KDim) + + return dict( + z_rho_e=z_rho_e, + z_vn_avg=z_vn_avg, + ddqz_z_full_e=ddqz_z_full_e, + z_theta_v_e=z_theta_v_e, + mass_fl_e=mass_fl_e, + z_theta_v_fl_e=z_theta_v_fl_e, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_33.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_33.py new file mode 100644 index 0000000000..4932c10f43 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_33.py @@ -0,0 +1,42 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_33 import ( + mo_solve_nonhydro_stencil_33, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil33(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_33 + OUTPUTS = ("vn_traj", "mass_flx_me") + + @staticmethod + def reference(mesh, vn_traj: np.array, mass_flx_me: np.array, **kwargs) -> dict: + vn_traj = np.zeros_like(vn_traj) + mass_flx_me = np.zeros_like(mass_flx_me) + return dict(vn_traj=vn_traj, mass_flx_me=mass_flx_me) + + @pytest.fixture + def input_data(self, mesh): + vn_traj = zero_field(mesh, EdgeDim, KDim) + mass_flx_me = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn_traj=vn_traj, + mass_flx_me=mass_flx_me, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_34.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_34.py new file mode 100644 index 0000000000..fabcf80228 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_34.py @@ -0,0 +1,56 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_34 import ( + mo_solve_nonhydro_stencil_34, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil34(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_34 + OUTPUTS = ("vn_traj", "mass_flx_me") + + @staticmethod + def reference( + mesh, + z_vn_avg: np.array, + mass_fl_e: np.array, + vn_traj: np.array, + mass_flx_me: np.array, + r_nsubsteps, + **kwargs, + ) -> dict: + vn_traj = vn_traj + r_nsubsteps * z_vn_avg + mass_flx_me = mass_flx_me + r_nsubsteps * mass_fl_e + return dict(vn_traj=vn_traj, mass_flx_me=mass_flx_me) + + @pytest.fixture + def input_data(self, mesh): + mass_fl_e = random_field(mesh, EdgeDim, KDim) + mass_flx_me = random_field(mesh, EdgeDim, KDim) + z_vn_avg = random_field(mesh, EdgeDim, KDim) + vn_traj = random_field(mesh, EdgeDim, KDim) + r_nsubsteps = 9.0 + + return dict( + z_vn_avg=z_vn_avg, + mass_fl_e=mass_fl_e, + vn_traj=vn_traj, + mass_flx_me=mass_flx_me, + r_nsubsteps=r_nsubsteps, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_35.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_35.py new file mode 100644 index 0000000000..3b087e61ef --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_35.py @@ -0,0 +1,58 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_35 import ( + mo_solve_nonhydro_stencil_35, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil35(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_35 + OUTPUTS = ("z_w_concorr_me",) + + @staticmethod + def reference( + mesh, + vn: np.array, + ddxn_z_full: np.array, + ddxt_z_full: np.array, + vt: np.array, + **kwargs, + ) -> dict: + z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full + return dict(z_w_concorr_me=z_w_concorr_me) + + @pytest.fixture + def input_data(self, mesh): + vn = random_field(mesh, EdgeDim, KDim) + ddxn_z_full = random_field(mesh, EdgeDim, KDim) + ddxt_z_full = random_field(mesh, EdgeDim, KDim) + vt = random_field(mesh, EdgeDim, KDim) + z_w_concorr_me = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn=vn, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + vt=vt, + z_w_concorr_me=z_w_concorr_me, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py similarity index 87% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_36.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py index 523c833dd1..86ded8421e 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_36.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_36 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_36 import ( mo_solve_nonhydro_stencil_36, ) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_36_numpy( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_37.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_37.py new file mode 100644 index 0000000000..8fdb06e2aa --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_37.py @@ -0,0 +1,53 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_37 import ( + mo_solve_nonhydro_stencil_37, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil37(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_37 + OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") + + @staticmethod + def reference(mesh, vn: np.array, vt: np.array, **kwargs) -> dict: + vn_ie = vn + z_vt_ie = vt + z_kin_hor_e = 0.5 * (pow(vn, 2) + pow(vt, 2)) + return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) + + @pytest.fixture + def input_data(self, mesh): + vt = random_field(mesh, EdgeDim, KDim) + vn = random_field(mesh, EdgeDim, KDim) + vn_ie = zero_field(mesh, EdgeDim, KDim) + z_kin_hor_e = zero_field(mesh, EdgeDim, KDim) + z_vt_ie = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn=vn, + vt=vt, + vn_ie=vn_ie, + z_vt_ie=z_vt_ie, + z_kin_hor_e=z_kin_hor_e, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py similarity index 84% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_38.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py index baaabf11f5..486f6f3289 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_38.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_38 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_38 import ( mo_solve_nonhydro_stencil_38, ) -from icon4py.common.dimension import EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_38_numpy( diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py similarity index 86% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index 1fbef1a76f..6e88bd7443 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_39 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_39 import ( mo_solve_nonhydro_stencil_39, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_39_numpy( diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py similarity index 88% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index ab0fce797d..082951fbd4 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_40 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_40 import ( mo_solve_nonhydro_stencil_40, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_40_numpy( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_41.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_41.py new file mode 100644 index 0000000000..62b09c102f --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_41.py @@ -0,0 +1,59 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_41 import ( + mo_solve_nonhydro_stencil_41, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil41(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_41 + OUTPUTS = ("z_flxdiv_mass", "z_flxdiv_theta") + + @staticmethod + def reference( + mesh, + geofac_div: np.array, + mass_fl_e: np.array, + z_theta_v_fl_e: np.array, + **kwargs, + ) -> tuple[np.array]: + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_flxdiv_mass = np.sum(geofac_div * mass_fl_e[mesh.c2e], axis=1) + z_flxdiv_theta = np.sum(geofac_div * z_theta_v_fl_e[mesh.c2e], axis=1) + return dict(z_flxdiv_mass=z_flxdiv_mass, z_flxdiv_theta=z_flxdiv_theta) + + @pytest.fixture + def input_data(self, mesh): + geofac_div = random_field(mesh, CellDim, C2EDim) + z_theta_v_fl_e = random_field(mesh, EdgeDim, KDim) + z_flxdiv_theta = zero_field(mesh, CellDim, KDim) + mass_fl_e = random_field(mesh, EdgeDim, KDim) + z_flxdiv_mass = zero_field(mesh, CellDim, KDim) + + return dict( + geofac_div=geofac_div, + mass_fl_e=mass_fl_e, + z_theta_v_fl_e=z_theta_v_fl_e, + z_flxdiv_mass=z_flxdiv_mass, + z_flxdiv_theta=z_flxdiv_theta, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py new file mode 100644 index 0000000000..7e0ea0bdd9 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py @@ -0,0 +1,85 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_42 import ( + mo_solve_nonhydro_stencil_42, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil42(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_42 + OUTPUTS = ("z_w_expl", "z_contr_w_fl_l") + + @staticmethod + def reference( + mesh, + w_nnow: np.array, + ddt_w_adv_ntl1: np.array, + ddt_w_adv_ntl2: np.array, + z_th_ddz_exner_c: np.array, + rho_ic: np.array, + w_concorr_c: np.array, + vwind_expl_wgt: np.array, + dtime, + wgt_nnow_vel, + wgt_nnew_vel, + cpd, + **kwargs, + ) -> tuple[np.array]: + z_w_expl = w_nnow + dtime * ( + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c + ) + vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) + z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) + return dict(z_w_expl=z_w_expl, z_contr_w_fl_l=z_contr_w_fl_l) + + @pytest.fixture + def input_data(self, mesh): + w_nnow = random_field(mesh, CellDim, KDim) + ddt_w_adv_ntl1 = random_field(mesh, CellDim, KDim) + ddt_w_adv_ntl2 = random_field(mesh, CellDim, KDim) + z_th_ddz_exner_c = random_field(mesh, CellDim, KDim) + z_w_expl = zero_field(mesh, CellDim, KDim) + rho_ic = random_field(mesh, CellDim, KDim) + w_concorr_c = random_field(mesh, CellDim, KDim) + vwind_expl_wgt = random_field(mesh, CellDim) + z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) + dtime = 5.0 + wgt_nnow_vel = 8.0 + wgt_nnew_vel = 9.0 + cpd = 10.0 + + return dict( + z_w_expl=z_w_expl, + w_nnow=w_nnow, + ddt_w_adv_ntl1=ddt_w_adv_ntl1, + ddt_w_adv_ntl2=ddt_w_adv_ntl2, + z_th_ddz_exner_c=z_th_ddz_exner_c, + z_contr_w_fl_l=z_contr_w_fl_l, + rho_ic=rho_ic, + w_concorr_c=w_concorr_c, + vwind_expl_wgt=vwind_expl_wgt, + dtime=dtime, + wgt_nnow_vel=wgt_nnow_vel, + wgt_nnew_vel=wgt_nnew_vel, + cpd=cpd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_43.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_43.py new file mode 100644 index 0000000000..4e2fb37562 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_43.py @@ -0,0 +1,74 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_43 import ( + mo_solve_nonhydro_stencil_43, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil43(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_43 + OUTPUTS = ("z_w_expl", "z_contr_w_fl_l") + + @staticmethod + def reference( + mesh, + w_nnow: np.array, + ddt_w_adv_ntl1: np.array, + z_th_ddz_exner_c: np.array, + rho_ic: np.array, + w_concorr_c: np.array, + vwind_expl_wgt: np.array, + dtime: float, + cpd: float, + **kwargs, + ) -> dict: + vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, -1) + z_w_expl = w_nnow + dtime * (ddt_w_adv_ntl1 - cpd * z_th_ddz_exner_c) + z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) + return dict(z_w_expl=z_w_expl, z_contr_w_fl_l=z_contr_w_fl_l) + + @pytest.fixture + def input_data(self, mesh): + w_nnow = random_field(mesh, CellDim, KDim) + ddt_w_adv_ntl1 = random_field(mesh, CellDim, KDim) + z_th_ddz_exner_c = random_field(mesh, CellDim, KDim) + z_w_expl = zero_field(mesh, CellDim, KDim) + rho_ic = random_field(mesh, CellDim, KDim) + w_concorr_c = random_field(mesh, CellDim, KDim) + vwind_expl_wgt = random_field(mesh, CellDim) + z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) + dtime = 5.0 + cpd = 10.0 + + return dict( + z_w_expl=z_w_expl, + w_nnow=w_nnow, + ddt_w_adv_ntl1=ddt_w_adv_ntl1, + z_th_ddz_exner_c=z_th_ddz_exner_c, + z_contr_w_fl_l=z_contr_w_fl_l, + rho_ic=rho_ic, + w_concorr_c=w_concorr_c, + vwind_expl_wgt=vwind_expl_wgt, + dtime=dtime, + cpd=cpd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py new file mode 100644 index 0000000000..2660d320ef --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py @@ -0,0 +1,81 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_44 import ( + mo_solve_nonhydro_stencil_44, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil44(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_44 + OUTPUTS = ("z_beta", "z_alpha") + + @staticmethod + def reference( + mesh, + exner_nnow: np.array, + rho_nnow: np.array, + theta_v_nnow: np.array, + inv_ddqz_z_full: np.array, + vwind_impl_wgt: np.array, + theta_v_ic: np.array, + rho_ic: np.array, + dtime, + rd, + cvd, + **kwargs, + ) -> dict: + z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full + + vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) + z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic + return dict(z_beta=z_beta, z_alpha=z_alpha) + + @pytest.fixture + def input_data(self, mesh): + exner_nnow = random_field(mesh, CellDim, KDim) + rho_nnow = random_field(mesh, CellDim, KDim) + theta_v_nnow = random_field(mesh, CellDim, KDim) + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + vwind_impl_wgt = random_field(mesh, CellDim) + theta_v_ic = random_field(mesh, CellDim, KDim) + rho_ic = random_field(mesh, CellDim, KDim) + z_alpha = zero_field(mesh, CellDim, KDim) + z_beta = zero_field(mesh, CellDim, KDim) + dtime = 10.0 + rd = 5.0 + cvd = 3.0 + + return dict( + z_beta=z_beta, + exner_nnow=exner_nnow, + rho_nnow=rho_nnow, + theta_v_nnow=theta_v_nnow, + inv_ddqz_z_full=inv_ddqz_z_full, + z_alpha=z_alpha, + vwind_impl_wgt=vwind_impl_wgt, + theta_v_ic=theta_v_ic, + rho_ic=rho_ic, + dtime=dtime, + rd=rd, + cvd=cvd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_45.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_45.py new file mode 100644 index 0000000000..8e1b453dd4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_45.py @@ -0,0 +1,39 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_45 import ( + mo_solve_nonhydro_stencil_45, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil45(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_45 + OUTPUTS = ("z_alpha",) + + @staticmethod + def reference(mesh, z_alpha: np.array, **kwargs) -> dict: + z_alpha = np.zeros_like(z_alpha) + return dict(z_alpha=z_alpha) + + @pytest.fixture + def input_data(self, mesh): + z_alpha = zero_field(mesh, CellDim, KDim) + + return dict( + z_alpha=z_alpha, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_46.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_46.py new file mode 100644 index 0000000000..cc9b1ecb5d --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_46.py @@ -0,0 +1,42 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_46 import ( + mo_solve_nonhydro_stencil_46, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil46(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_46 + OUTPUTS = ("w_nnew", "z_contr_w_fl_l") + + @staticmethod + def reference(mesh, w_nnew: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: + w_nnew = np.zeros_like(w_nnew) + z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) + return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) + + @pytest.fixture + def input_data(self, mesh): + z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) + w_nnew = zero_field(mesh, CellDim, KDim) + + return dict( + w_nnew=w_nnew, + z_contr_w_fl_l=z_contr_w_fl_l, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py new file mode 100644 index 0000000000..60253bc78e --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_47 import ( + mo_solve_nonhydro_stencil_47, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil47(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_47 + OUTPUTS = ("w_nnew", "z_contr_w_fl_l") + + @staticmethod + def reference(mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: + w_nnew = w_concorr_c + z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) + return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) + + @pytest.fixture + def input_data(self, mesh): + w_concorr_c = random_field(mesh, CellDim, KDim) + z_contr_w_fl_l = zero_field(mesh, CellDim, KDim) + w_nnew = zero_field(mesh, CellDim, KDim) + + return dict( + w_nnew=w_nnew, + z_contr_w_fl_l=z_contr_w_fl_l, + w_concorr_c=w_concorr_c, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_48.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_48.py new file mode 100644 index 0000000000..6654f3177d --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_48.py @@ -0,0 +1,92 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_48 import ( + mo_solve_nonhydro_stencil_48, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil48(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_48 + OUTPUTS = ("z_rho_expl", "z_exner_expl") + + @staticmethod + def reference( + mesh, + rho_nnow: np.array, + inv_ddqz_z_full: np.array, + z_flxdiv_mass: np.array, + z_contr_w_fl_l: np.array, + exner_pr: np.array, + z_beta: np.array, + z_flxdiv_theta: np.array, + theta_v_ic: np.array, + ddt_exner_phy: np.array, + dtime, + **kwargs, + ) -> dict: + z_rho_expl = rho_nnow - dtime * inv_ddqz_z_full * ( + z_flxdiv_mass + z_contr_w_fl_l[:, :-1] - z_contr_w_fl_l[:, 1:] + ) + + z_exner_expl = ( + exner_pr + - z_beta + * ( + z_flxdiv_theta + + (theta_v_ic * z_contr_w_fl_l)[:, :-1] + - (theta_v_ic * z_contr_w_fl_l)[:, 1:] + ) + + dtime * ddt_exner_phy + ) + return dict(z_rho_expl=z_rho_expl, z_exner_expl=z_exner_expl) + + @pytest.fixture + def input_data(self, mesh): + dtime = 1.0 + rho_nnow = random_field(mesh, CellDim, KDim) + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + z_flxdiv_mass = random_field(mesh, CellDim, KDim) + z_contr_w_fl_l = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + exner_pr = random_field(mesh, CellDim, KDim) + z_beta = random_field(mesh, CellDim, KDim) + z_flxdiv_theta = random_field(mesh, CellDim, KDim) + theta_v_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + ddt_exner_phy = random_field(mesh, CellDim, KDim) + + z_rho_expl = zero_field(mesh, CellDim, KDim) + z_exner_expl = zero_field(mesh, CellDim, KDim) + + return dict( + z_rho_expl=z_rho_expl, + z_exner_expl=z_exner_expl, + rho_nnow=rho_nnow, + inv_ddqz_z_full=inv_ddqz_z_full, + z_flxdiv_mass=z_flxdiv_mass, + z_contr_w_fl_l=z_contr_w_fl_l, + exner_pr=exner_pr, + z_beta=z_beta, + z_flxdiv_theta=z_flxdiv_theta, + theta_v_ic=theta_v_ic, + ddt_exner_phy=ddt_exner_phy, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_49.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_49.py new file mode 100644 index 0000000000..98c5f6e8ca --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_49.py @@ -0,0 +1,92 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_49 import ( + mo_solve_nonhydro_stencil_49, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil49(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_49 + OUTPUTS = ("z_rho_expl", "z_exner_expl") + + @staticmethod + def reference( + mesh, + rho_nnow: np.array, + inv_ddqz_z_full: np.array, + z_flxdiv_mass: np.array, + z_contr_w_fl_l: np.array, + exner_pr: np.array, + z_beta: np.array, + z_flxdiv_theta: np.array, + theta_v_ic: np.array, + ddt_exner_phy: np.array, + dtime, + **kwargs, + ) -> dict: + z_rho_expl = rho_nnow - dtime * inv_ddqz_z_full * ( + z_flxdiv_mass + z_contr_w_fl_l[:, :-1] - z_contr_w_fl_l[:, 1:] + ) + z_exner_expl = ( + exner_pr + - z_beta + * ( + z_flxdiv_theta + + (theta_v_ic * z_contr_w_fl_l)[:, :-1] + - (theta_v_ic * z_contr_w_fl_l)[:, 1:] + ) + + dtime * ddt_exner_phy + ) + return dict(z_rho_expl=z_rho_expl, z_exner_expl=z_exner_expl) + + @pytest.fixture + def input_data(self, mesh): + dtime = 7.0 + + rho_nnow = random_field(mesh, CellDim, KDim) + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + z_flxdiv_mass = random_field(mesh, CellDim, KDim) + z_contr_w_fl_l = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + exner_pr = random_field(mesh, CellDim, KDim) + z_beta = random_field(mesh, CellDim, KDim) + z_flxdiv_theta = random_field(mesh, CellDim, KDim) + theta_v_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + ddt_exner_phy = random_field(mesh, CellDim, KDim) + + z_rho_expl = zero_field(mesh, CellDim, KDim) + z_exner_expl = zero_field(mesh, CellDim, KDim) + + return dict( + z_rho_expl=z_rho_expl, + z_exner_expl=z_exner_expl, + rho_nnow=rho_nnow, + inv_ddqz_z_full=inv_ddqz_z_full, + z_flxdiv_mass=z_flxdiv_mass, + z_contr_w_fl_l=z_contr_w_fl_l, + exner_pr=exner_pr, + z_beta=z_beta, + z_flxdiv_theta=z_flxdiv_theta, + theta_v_ic=theta_v_ic, + ddt_exner_phy=ddt_exner_phy, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_50.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_50.py new file mode 100644 index 0000000000..f0da6ef58b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_50.py @@ -0,0 +1,56 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_50 import ( + mo_solve_nonhydro_stencil_50, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil50(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_50 + OUTPUTS = ("z_rho_expl", "z_exner_expl") + + @staticmethod + def reference( + mesh, + z_rho_expl: np.array, + rho_incr: np.array, + z_exner_expl: np.array, + exner_incr: np.array, + iau_wgt_dyn, + **kwargs, + ) -> dict: + z_rho_expl = z_rho_expl + iau_wgt_dyn * rho_incr + z_exner_expl = z_exner_expl + iau_wgt_dyn * exner_incr + return dict(z_rho_expl=z_rho_expl, z_exner_expl=z_exner_expl) + + @pytest.fixture + def input_data(self, mesh): + z_exner_expl = random_field(mesh, CellDim, KDim) + exner_incr = random_field(mesh, CellDim, KDim) + z_rho_expl = random_field(mesh, CellDim, KDim) + rho_incr = random_field(mesh, CellDim, KDim) + iau_wgt_dyn = 8.0 + + return dict( + z_rho_expl=z_rho_expl, + z_exner_expl=z_exner_expl, + rho_incr=rho_incr, + exner_incr=exner_incr, + iau_wgt_dyn=iau_wgt_dyn, + ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py similarity index 89% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_51.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py index 0140a842f2..f7e12c96f2 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py @@ -13,20 +13,18 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_51 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_51 import ( mo_solve_nonhydro_stencil_51, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_51_z_q_numpy( z_c: np.array, z_b: np.array, ) -> np.array: - return -z_c / z_b @@ -59,9 +57,7 @@ def mo_solve_nonhydro_stencil_51_numpy( z_c = -z_gamma * z_beta * z_alpha_k_plus_1 z_b = 1.0 + z_gamma * z_alpha[:, :-1] * (z_beta_k_minus_1 + z_beta) z_q = mo_solve_nonhydro_stencil_51_z_q_numpy(z_c, z_b) - w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy( - z_gamma, z_b, z_w_expl, z_exner_expl - ) + w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy(z_gamma, z_b, z_w_expl, z_exner_expl) return z_q, w_nnew diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py similarity index 87% rename from atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_52.py rename to model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py index c2ab678402..6c0cc28e0d 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py @@ -14,13 +14,12 @@ import numpy as np from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_52 import ( +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_52 import ( mo_solve_nonhydro_stencil_52, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_solve_nonhydro_stencil_52_numpy( @@ -50,15 +49,11 @@ def mo_solve_nonhydro_stencil_52_numpy( for k in range(1, k_size): z_a[:, k] = -z_gamma[:, k] * z_beta[:, k - 1] * z_alpha[:, k - 1] z_c[:, k] = -z_gamma[:, k] * z_beta[:, k] * z_alpha[:, k + 1] - z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * ( - z_beta[:, k - 1] + z_beta[:, k] - ) + z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * (z_beta[:, k - 1] + z_beta[:, k]) z_g[:, k] = 1.0 / (z_b[:, k] + z_a[:, k] * z_q[:, k - 1]) z_q[:, k] = -z_c[:, k] * z_g[:, k] - w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * ( - z_exner_expl[:, k - 1] - z_exner_expl[:, k] - ) + w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * (z_exner_expl[:, k - 1] - z_exner_expl[:, k]) w[:, k] = (w[:, k] - z_a[:, k] * w[:, k - 1]) * z_g[:, k] return z_q, w diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_53.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_53.py new file mode 100644 index 0000000000..a422cbcda7 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_53.py @@ -0,0 +1,43 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_53 import ( + mo_solve_nonhydro_stencil_53, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil53(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_53 + OUTPUTS = ("w",) + + @staticmethod + def reference(mesh, z_q: np.array, w: np.array, **kwargs) -> np.array: + w_new = np.zeros_like(w) + last_k_level = w.shape[1] - 1 + + w_new[:, last_k_level] = w[:, last_k_level] + for k in reversed(range(1, last_k_level)): + w_new[:, k] = w[:, k] + w_new[:, k + 1] * z_q[:, k] + w_new[:, 0] = w[:, 0] + return dict(w=w_new) + + @pytest.fixture + def input_data(self, mesh): + z_q = random_field(mesh, CellDim, KDim) + w = random_field(mesh, CellDim, KDim) + return dict(z_q=z_q, w=w) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py new file mode 100644 index 0000000000..eeb17c2cce --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py @@ -0,0 +1,45 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_54 import ( + mo_solve_nonhydro_stencil_54, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil54(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_54 + OUTPUTS = ("w",) + + @staticmethod + def reference(mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -> np.array: + z_raylfac = np.expand_dims(z_raylfac, axis=0) + w_1 = np.expand_dims(w_1, axis=-1) + w = z_raylfac * w + (1.0 - z_raylfac) * w_1 + return dict(w=w) + + @pytest.fixture + def input_data(self, mesh): + z_raylfac = random_field(mesh, KDim) + w_1 = random_field(mesh, CellDim) + w = random_field(mesh, CellDim, KDim) + + return dict( + z_raylfac=z_raylfac, + w_1=w_1, + w=w, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py new file mode 100644 index 0000000000..1d9e126493 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py @@ -0,0 +1,107 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_55 import ( + mo_solve_nonhydro_stencil_55, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil55(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_55 + OUTPUTS = ("rho_new", "exner_new", "theta_v_new") + + @staticmethod + def reference( + mesh, + z_rho_expl: np.array, + vwind_impl_wgt: np.array, + inv_ddqz_z_full: np.array, + rho_ic: np.array, + w: np.array, + z_exner_expl: np.array, + exner_ref_mc: np.array, + z_alpha: np.array, + z_beta: np.array, + rho_now: np.array, + theta_v_now: np.array, + exner_now: np.array, + dtime, + cvd_o_rd, + **kwargs, + ) -> dict: + rho_ic_offset_1 = rho_ic[:, 1:] + w_offset_0 = w[:, :-1] + w_offset_1 = w[:, 1:] + z_alpha_offset_1 = z_alpha[:, 1:] + vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=1) + rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( + rho_ic[:, :-1] * w_offset_0 - rho_ic_offset_1 * w_offset_1 + ) + exner_new = ( + z_exner_expl + + exner_ref_mc + - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) + ) + theta_v_new = ( + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new + ) + return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) + + @pytest.fixture + def input_data(self, mesh): + z_rho_expl = random_field(mesh, CellDim, KDim) + vwind_impl_wgt = random_field(mesh, CellDim) + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + rho_ic = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + w = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + z_exner_expl = random_field(mesh, CellDim, KDim) + exner_ref_mc = random_field(mesh, CellDim, KDim) + z_alpha = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + z_beta = random_field(mesh, CellDim, KDim) + rho_now = random_field(mesh, CellDim, KDim) + theta_v_now = random_field(mesh, CellDim, KDim) + exner_now = random_field(mesh, CellDim, KDim) + rho_new = zero_field(mesh, CellDim, KDim) + exner_new = zero_field(mesh, CellDim, KDim) + theta_v_new = zero_field(mesh, CellDim, KDim) + dtime = 5.0 + cvd_o_rd = 9.0 + + return dict( + z_rho_expl=z_rho_expl, + vwind_impl_wgt=vwind_impl_wgt, + inv_ddqz_z_full=inv_ddqz_z_full, + rho_ic=rho_ic, + w=w, + z_exner_expl=z_exner_expl, + exner_ref_mc=exner_ref_mc, + z_alpha=z_alpha, + z_beta=z_beta, + rho_now=rho_now, + theta_v_now=theta_v_now, + exner_now=exner_now, + rho_new=rho_new, + exner_new=exner_new, + theta_v_new=theta_v_new, + dtime=dtime, + cvd_o_rd=cvd_o_rd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_56_63.py new file mode 100644 index 0000000000..5b8d1b4757 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_56_63.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_56_63 import ( + mo_solve_nonhydro_stencil_56_63, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil5663(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_56_63 + OUTPUTS = ("z_dwdz_dd",) + + @staticmethod + def reference( + mesh, inv_ddqz_z_full: np.array, w: np.array, w_concorr_c: np.array, **kwargs + ) -> np.array: + z_dwdz_dd = inv_ddqz_z_full * ( + (w[:, :-1] - w[:, 1:]) - (w_concorr_c[:, :-1] - w_concorr_c[:, 1:]) + ) + return dict(z_dwdz_dd=z_dwdz_dd) + + @pytest.fixture + def input_data(self, mesh): + inv_ddqz_z_full = random_field(mesh, CellDim, KDim) + w = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + w_concorr_c = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + z_dwdz_dd = random_field(mesh, CellDim, KDim) + + return dict( + inv_ddqz_z_full=inv_ddqz_z_full, + w=w, + w_concorr_c=w_concorr_c, + z_dwdz_dd=z_dwdz_dd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_57.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_57.py new file mode 100644 index 0000000000..b6f80122b1 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_57.py @@ -0,0 +1,39 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_57 import ( + mo_solve_nonhydro_stencil_57, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil57(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_57 + OUTPUTS = ("mass_flx_ic",) + + @staticmethod + def reference(mesh, mass_flx_ic: np.array, **kwargs) -> dict: + mass_flx_ic = np.zeros_like(mass_flx_ic) + return dict(mass_flx_ic=mass_flx_ic) + + @pytest.fixture + def input_data(self, mesh): + mass_flx_ic = zero_field(mesh, CellDim, KDim) + + return dict( + mass_flx_ic=mass_flx_ic, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py new file mode 100644 index 0000000000..dff001052b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py @@ -0,0 +1,59 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_58 import ( + mo_solve_nonhydro_stencil_58, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil58(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_58 + OUTPUTS = ("mass_flx_ic",) + + @staticmethod + def reference( + mesh, + z_contr_w_fl_l: np.array, + rho_ic: np.array, + vwind_impl_wgt: np.array, + w: np.array, + mass_flx_ic: np.array, + r_nsubsteps, + **kwargs, + ) -> dict: + vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) + return dict(mass_flx_ic=mass_flx_ic) + + @pytest.fixture + def input_data(self, mesh): + z_contr_w_fl_l = random_field(mesh, CellDim, KDim) + rho_ic = random_field(mesh, CellDim, KDim) + vwind_impl_wgt = random_field(mesh, CellDim) + w = random_field(mesh, CellDim, KDim) + mass_flx_ic = random_field(mesh, CellDim, KDim) + r_nsubsteps = 7.0 + + return dict( + z_contr_w_fl_l=z_contr_w_fl_l, + rho_ic=rho_ic, + vwind_impl_wgt=vwind_impl_wgt, + w=w, + mass_flx_ic=mass_flx_ic, + r_nsubsteps=r_nsubsteps, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_59.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_59.py new file mode 100644 index 0000000000..6fead774be --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_59.py @@ -0,0 +1,45 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_59 import ( + mo_solve_nonhydro_stencil_59, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil59(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_59 + OUTPUTS = ("exner_dyn_incr",) + + @staticmethod + def reference(mesh, exner: np.array, **kwargs) -> dict: + exner_dyn_incr = exner + return dict(exner_dyn_incr=exner_dyn_incr) + + @pytest.fixture + def input_data(self, mesh): + exner = random_field(mesh, CellDim, KDim) + exner_dyn_incr = zero_field(mesh, CellDim, KDim) + + return dict( + exner=exner, + exner_dyn_incr=exner_dyn_incr, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py new file mode 100644 index 0000000000..634446cb33 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py @@ -0,0 +1,54 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_60 import ( + mo_solve_nonhydro_stencil_60, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil60(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_60 + OUTPUTS = ("exner_dyn_incr",) + + @staticmethod + def reference( + mesh, + exner: np.array, + ddt_exner_phy: np.array, + exner_dyn_incr: np.array, + ndyn_substeps_var: float, + dtime: float, + **kwargs, + ) -> np.array: + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) + return dict(exner_dyn_incr=exner_dyn_incr) + + @pytest.fixture + def input_data(self, mesh): + ndyn_substeps_var, dtime = 10.0, 12.0 + exner = random_field(mesh, CellDim, KDim) + ddt_exner_phy = random_field(mesh, CellDim, KDim) + exner_dyn_incr = random_field(mesh, CellDim, KDim) + + return dict( + exner=exner, + ddt_exner_phy=ddt_exner_phy, + exner_dyn_incr=exner_dyn_incr, + ndyn_substeps_var=ndyn_substeps_var, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_61.py new file mode 100644 index 0000000000..ec105b2996 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_61.py @@ -0,0 +1,73 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_61 import ( + mo_solve_nonhydro_stencil_61, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil61(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_61 + OUTPUTS = ("rho_new", "exner_new", "w_new") + + @staticmethod + def reference( + mesh, + rho_now: np.array, + grf_tend_rho: np.array, + theta_v_now: np.array, + grf_tend_thv: np.array, + w_now: np.array, + grf_tend_w: np.array, + dtime, + **kwargs, + ) -> tuple[np.array]: + rho_new = rho_now + dtime * grf_tend_rho + exner_new = theta_v_now + dtime * grf_tend_thv + w_new = w_now + dtime * grf_tend_w + return dict(rho_new=rho_new, exner_new=exner_new, w_new=w_new) + + @pytest.fixture + def input_data(self, mesh): + rho_now = random_field(mesh, CellDim, KDim) + grf_tend_rho = random_field(mesh, CellDim, KDim) + theta_v_now = random_field(mesh, CellDim, KDim) + grf_tend_thv = random_field(mesh, CellDim, KDim) + w_now = random_field(mesh, CellDim, KDim) + grf_tend_w = random_field(mesh, CellDim, KDim) + dtime = 5.0 + rho_new = zero_field(mesh, CellDim, KDim) + exner_new = zero_field(mesh, CellDim, KDim) + w_new = zero_field(mesh, CellDim, KDim) + + return dict( + rho_now=rho_now, + grf_tend_rho=grf_tend_rho, + theta_v_now=theta_v_now, + grf_tend_thv=grf_tend_thv, + w_now=w_now, + grf_tend_w=grf_tend_w, + rho_new=rho_new, + exner_new=exner_new, + w_new=w_new, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py new file mode 100644 index 0000000000..41ef7c9439 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py @@ -0,0 +1,49 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_62 import ( + mo_solve_nonhydro_stencil_62, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoSolveNonhydroStencil62(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_62 + OUTPUTS = ("w_new",) + + @staticmethod + def reference(mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs) -> np.array: + w_new = w_now + dtime * grf_tend_w + return dict(w_new=w_new) + + @pytest.fixture + def input_data(self, mesh): + dtime = 10.0 + w_now = random_field(mesh, CellDim, KDim) + grf_tend_w = random_field(mesh, CellDim, KDim) + w_new = zero_field(mesh, CellDim, KDim) + + return dict( + w_now=w_now, + grf_tend_w=grf_tend_w, + w_new=w_new, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_64.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_64.py new file mode 100644 index 0000000000..be211ddcdf --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_64.py @@ -0,0 +1,39 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_64 import ( + mo_solve_nonhydro_stencil_64, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, zero_field + + +class TestMoSolveNonhydroStencil64(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_64 + OUTPUTS = ("mass_flx_ic",) + + @staticmethod + def reference(mesh, mass_flx_ic: np.array, **kwargs) -> dict: + mass_flx_ic = np.zeros_like(mass_flx_ic) + return dict(mass_flx_ic=mass_flx_ic) + + @pytest.fixture + def input_data(self, mesh): + mass_flx_ic = zero_field(mesh, CellDim, KDim) + + return dict( + mass_flx_ic=mass_flx_ic, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py new file mode 100644 index 0000000000..b5b5b9c0c0 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py @@ -0,0 +1,68 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_65 import ( + mo_solve_nonhydro_stencil_65, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil65(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_65 + OUTPUTS = ("mass_flx_ic",) + + @staticmethod + def reference( + mesh, + rho_ic: np.array, + vwind_expl_wgt: np.array, + vwind_impl_wgt: np.array, + w_now: np.array, + w_new: np.array, + w_concorr_c: np.array, + mass_flx_ic: np.array, + r_nsubsteps: float, + **kwargs, + ) -> np.array: + vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) + vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) + mass_flx_ic = mass_flx_ic + ( + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + ) + return dict(mass_flx_ic=mass_flx_ic) + + @pytest.fixture + def input_data(self, mesh): + r_nsubsteps = 10.0 + rho_ic = random_field(mesh, CellDim, KDim) + vwind_expl_wgt = random_field(mesh, CellDim) + vwind_impl_wgt = random_field(mesh, CellDim) + w_now = random_field(mesh, CellDim, KDim) + w_new = random_field(mesh, CellDim, KDim) + w_concorr_c = random_field(mesh, CellDim, KDim) + mass_flx_ic = random_field(mesh, CellDim, KDim) + + return dict( + rho_ic=rho_ic, + vwind_expl_wgt=vwind_expl_wgt, + vwind_impl_wgt=vwind_impl_wgt, + w_now=w_now, + w_new=w_new, + w_concorr_c=w_concorr_c, + mass_flx_ic=mass_flx_ic, + r_nsubsteps=r_nsubsteps, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_66.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_66.py new file mode 100644 index 0000000000..cc044558b4 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_66.py @@ -0,0 +1,68 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_66 import ( + mo_solve_nonhydro_stencil_66, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, +) + + +class TestMoSolveNonhydroStencil66(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_66 + OUTPUTS = ("theta_v", "exner") + + @staticmethod + def reference( + mesh, + bdy_halo_c: np.array, + rho: np.array, + theta_v: np.array, + exner: np.array, + rd_o_cvd: float, + rd_o_p0ref: float, + **kwargs, + ) -> dict: + bdy_halo_c = np.expand_dims(bdy_halo_c, axis=-1) + + theta_v = np.where(bdy_halo_c == 1, exner, theta_v) + exner = np.where( + bdy_halo_c == 1, np.exp(rd_o_cvd * np.log(rd_o_p0ref * rho * exner)), exner + ) + + return dict(theta_v=theta_v, exner=exner) + + @pytest.fixture + def input_data(self, mesh): + rd_o_cvd = 10.0 + rd_o_p0ref = 20.0 + bdy_halo_c = random_mask(mesh, CellDim) + exner = random_field(mesh, CellDim, KDim, low=1, high=2) + rho = random_field(mesh, CellDim, KDim, low=1, high=2) + theta_v = random_field(mesh, CellDim, KDim, low=1, high=2) + + return dict( + bdy_halo_c=bdy_halo_c, + rho=rho, + theta_v=theta_v, + exner=exner, + rd_o_cvd=rd_o_cvd, + rd_o_p0ref=rd_o_p0ref, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_67.py new file mode 100644 index 0000000000..eae749cd03 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_67.py @@ -0,0 +1,55 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_67 import ( + mo_solve_nonhydro_stencil_67, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoSolveNonhydroStencil67(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_67 + OUTPUTS = ("theta_v", "exner") + + @staticmethod + def reference( + mesh, + rho: np.array, + exner: np.array, + rd_o_cvd: float, + rd_o_p0ref: float, + **kwargs, + ) -> dict: + theta_v = np.copy(exner) + exner = np.exp(rd_o_cvd * np.log(rd_o_p0ref * rho * theta_v)) + return dict(theta_v=theta_v, exner=exner) + + @pytest.fixture + def input_data(self, mesh): + rd_o_cvd = 10.0 + rd_o_p0ref = 20.0 + rho = random_field(mesh, CellDim, KDim, low=1, high=2) + theta_v = random_field(mesh, CellDim, KDim, low=1, high=2) + exner = random_field(mesh, CellDim, KDim, low=1, high=2) + + return dict( + rho=rho, + theta_v=theta_v, + exner=exner, + rd_o_cvd=rd_o_cvd, + rd_o_p0ref=rd_o_p0ref, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py new file mode 100644 index 0000000000..dd54adca2b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py @@ -0,0 +1,74 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_68 import ( + mo_solve_nonhydro_stencil_68, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, +) + + +class TestMoSolveNonhydroStencil68(StencilTest): + PROGRAM = mo_solve_nonhydro_stencil_68 + OUTPUTS = ("theta_v_new",) + + @staticmethod + def reference( + mesh, + mask_prog_halo_c: np.array, + rho_now: np.array, + theta_v_now: np.array, + exner_new: np.array, + exner_now: np.array, + rho_new: np.array, + theta_v_new: np.array, + cvd_o_rd: float, + **kwargs, + ) -> dict: + mask_prog_halo_c = np.expand_dims(mask_prog_halo_c, axis=-1) + + theta_v_new = np.where( + mask_prog_halo_c, + rho_now * theta_v_now * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) / rho_new, + theta_v_new, + ) + return dict(theta_v_new=theta_v_new) + + @pytest.fixture + def input_data(self, mesh): + mask_prog_halo_c = random_mask(mesh, CellDim) + rho_now = random_field(mesh, CellDim, KDim) + theta_v_now = random_field(mesh, CellDim, KDim) + exner_new = random_field(mesh, CellDim, KDim) + exner_now = random_field(mesh, CellDim, KDim) + rho_new = random_field(mesh, CellDim, KDim) + theta_v_new = random_field(mesh, CellDim, KDim) + cvd_o_rd = 10.0 + + return dict( + mask_prog_halo_c=mask_prog_halo_c, + rho_now=rho_now, + theta_v_now=theta_v_now, + exner_new=exner_new, + exner_now=exner_now, + rho_new=rho_new, + theta_v_new=theta_v_new, + cvd_o_rd=cvd_o_rd, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_01.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_01.py new file mode 100644 index 0000000000..f72cd033b0 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_01.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_01 import ( + mo_velocity_advection_stencil_01, +) +from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil01(StencilTest): + PROGRAM = mo_velocity_advection_stencil_01 + OUTPUTS = ("vt",) + + @staticmethod + def reference(mesh, vn: np.array, rbf_vec_coeff_e: np.array, **kwargs) -> np.array: + rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) + vt = np.sum(vn[mesh.e2c2e] * rbf_vec_coeff_e, axis=1) + return dict(vt=vt) + + @pytest.fixture + def input_data(self, mesh): + vn = random_field(mesh, EdgeDim, KDim) + rbf_vec_coeff_e = random_field(mesh, EdgeDim, E2C2EDim) + vt = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn=vn, + rbf_vec_coeff_e=rbf_vec_coeff_e, + vt=vt, + ) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py similarity index 79% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_02.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index daa3152c56..9bdc11bf27 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -13,26 +13,21 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_02 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_02 import ( mo_velocity_advection_stencil_02, ) -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - -def mo_velocity_advection_stencil_02_vn_ie_numpy( - wgtfac_e: np.array, vn: np.array -) -> np.array: +def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie -def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( - vn: np.array, vt: np.array -) -> np.array: +def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py similarity index 70% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_03.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py index 293a22cfbe..f32c87b1b1 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py @@ -13,18 +13,15 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_03 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_03 import ( mo_velocity_advection_stencil_03, ) -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - -def mo_velocity_advection_stencil_03_numpy( - wgtfac_e: np.array, vt: np.array -) -> np.array: +def mo_velocity_advection_stencil_03_numpy(wgtfac_e: np.array, vt: np.array) -> np.array: vt_k_minus_1 = np.roll(vt, shift=1, axis=1) z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 @@ -39,9 +36,7 @@ def test_mo_velocity_advection_stencil_03(): z_vt_ie = zero_field(mesh, EdgeDim, KDim) - z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy( - np.asarray(wgtfac_e), np.asarray(vt) - ) + z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy(np.asarray(wgtfac_e), np.asarray(vt)) mo_velocity_advection_stencil_03( wgtfac_e, vt, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_04.py new file mode 100644 index 0000000000..cc124bf8bc --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_04.py @@ -0,0 +1,58 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import ( + mo_velocity_advection_stencil_04, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil04(StencilTest): + PROGRAM = mo_velocity_advection_stencil_04 + OUTPUTS = ("z_w_concorr_me",) + + @staticmethod + def reference( + mesh, + vn: np.array, + ddxn_z_full: np.array, + ddxt_z_full: np.array, + vt: np.array, + **kwargs, + ) -> np.array: + z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full + return dict(z_w_concorr_me=z_w_concorr_me) + + @pytest.fixture + def input_data(self, mesh): + vn = random_field(mesh, EdgeDim, KDim) + ddxn_z_full = random_field(mesh, EdgeDim, KDim) + ddxt_z_full = random_field(mesh, EdgeDim, KDim) + vt = random_field(mesh, EdgeDim, KDim) + z_w_concorr_me = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn=vn, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + vt=vt, + z_w_concorr_me=z_w_concorr_me, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_05.py new file mode 100644 index 0000000000..4ec8230b5f --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_05.py @@ -0,0 +1,54 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import ( + mo_velocity_advection_stencil_05, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil05(StencilTest): + PROGRAM = mo_velocity_advection_stencil_05 + OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") + + @staticmethod + def reference(mesh, vn: np.array, vt: np.array, **kwargs) -> dict: + vn_ie = vn + z_vt_ie = vt + z_kin_hor_e = 0.5 * ((vn * vn) + (vt * vt)) + return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) + + @pytest.fixture + def input_data(self, mesh): + vn = random_field(mesh, EdgeDim, KDim) + vt = random_field(mesh, EdgeDim, KDim) + + vn_ie = zero_field(mesh, EdgeDim, KDim) + z_vt_ie = zero_field(mesh, EdgeDim, KDim) + z_kin_hor_e = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn=vn, + vt=vt, + vn_ie=vn_ie, + z_vt_ie=z_vt_ie, + z_kin_hor_e=z_kin_hor_e, + ) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py similarity index 76% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_06.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py index 7a43be5c50..1acd23aeb1 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py @@ -13,18 +13,15 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_06 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import ( mo_velocity_advection_stencil_06, ) -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - -def mo_velocity_advection_stencil_06_numpy( - wgtfacq_e: np.array, vn: np.array -) -> np.array: +def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> np.array: vn_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_k_minus_2 = np.roll(vn, shift=2, axis=1) vn_k_minus_3 = np.roll(vn, shift=3, axis=1) @@ -48,9 +45,7 @@ def test_mo_velocity_advection_stencil_06(): vn_ie = zero_field(mesh, EdgeDim, KDim) - vn_ie_ref = mo_velocity_advection_stencil_06_numpy( - np.asarray(wgtfacq_e), np.asarray(vn) - ) + vn_ie_ref = mo_velocity_advection_stencil_06_numpy(np.asarray(wgtfacq_e), np.asarray(vn)) mo_velocity_advection_stencil_06( wgtfacq_e, vn, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_07.py new file mode 100644 index 0000000000..995babfd6f --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_07.py @@ -0,0 +1,80 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_07 import ( + mo_velocity_advection_stencil_07, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil07(StencilTest): + PROGRAM = mo_velocity_advection_stencil_07 + OUTPUTS = ("z_v_grad_w",) + + @staticmethod + def reference( + mesh, + vn_ie: np.array, + inv_dual_edge_length: np.array, + w: np.array, + z_vt_ie: np.array, + inv_primal_edge_length: np.array, + tangent_orientation: np.array, + z_w_v: np.array, + **kwargs, + ) -> np.array: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + + w_e2c = w[mesh.e2c] + z_w_v_e2v = z_w_v[mesh.e2v] + + red_w = w_e2c[:, 0] - w_e2c[:, 1] + red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] + + z_v_grad_w = ( + vn_ie * inv_dual_edge_length * red_w + + z_vt_ie * inv_primal_edge_length * tangent_orientation * red_z_w_v + ) + return dict(z_v_grad_w=z_v_grad_w) + + @pytest.fixture + def input_data(self, mesh): + vn_ie = random_field(mesh, EdgeDim, KDim) + inv_dual_edge_length = random_field(mesh, EdgeDim) + w = random_field(mesh, CellDim, KDim) + z_vt_ie = random_field(mesh, EdgeDim, KDim) + inv_primal_edge_length = random_field(mesh, EdgeDim) + tangent_orientation = random_field(mesh, EdgeDim) + z_w_v = random_field(mesh, VertexDim, KDim) + z_v_grad_w = zero_field(mesh, EdgeDim, KDim) + + return dict( + vn_ie=vn_ie, + inv_dual_edge_length=inv_dual_edge_length, + w=w, + z_vt_ie=z_vt_ie, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + z_w_v=z_w_v, + z_v_grad_w=z_v_grad_w, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py new file mode 100644 index 0000000000..192de3f3c3 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import ( + mo_velocity_advection_stencil_08, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil08(StencilTest): + PROGRAM = mo_velocity_advection_stencil_08 + OUTPUTS = ("z_ekinh",) + + @staticmethod + def reference(mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + z_ekinh = np.sum(z_kin_hor_e[mesh.c2e] * e_bln_c_s, axis=1) + return dict(z_ekinh=z_ekinh) + + @pytest.fixture + def input_data(self, mesh): + z_kin_hor_e = random_field(mesh, EdgeDim, KDim) + e_bln_c_s = random_field(mesh, CellDim, C2EDim) + z_ekinh = zero_field(mesh, CellDim, KDim) + + return dict( + z_kin_hor_e=z_kin_hor_e, + e_bln_c_s=e_bln_c_s, + z_ekinh=z_ekinh, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py new file mode 100644 index 0000000000..664b168561 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py @@ -0,0 +1,48 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import ( + mo_velocity_advection_stencil_09, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil09(StencilTest): + PROGRAM = mo_velocity_advection_stencil_09 + OUTPUTS = ("z_w_concorr_mc",) + + @staticmethod + def reference(mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + z_w_concorr_mc = np.sum(z_w_concorr_me[mesh.c2e] * e_bln_c_s, axis=1) + return dict(z_w_concorr_mc=z_w_concorr_mc) + + @pytest.fixture + def input_data(self, mesh): + z_w_concorr_me = random_field(mesh, EdgeDim, KDim) + e_bln_c_s = random_field(mesh, CellDim, C2EDim) + z_w_concorr_mc = zero_field(mesh, CellDim, KDim) + + return dict( + z_w_concorr_me=z_w_concorr_me, + e_bln_c_s=e_bln_c_s, + z_w_concorr_mc=z_w_concorr_mc, + ) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py similarity index 78% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_10.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index 0141eb4b69..33c9bfe847 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -13,22 +13,19 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_10 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_10 import ( mo_velocity_advection_stencil_10, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_velocity_advection_stencil_10_numpy( wgtfac_c: np.array, z_w_concorr_mc: np.array ) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = ( - wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 return w_concorr_c diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_11.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_11.py new file mode 100644 index 0000000000..dc16cec9cb --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_11.py @@ -0,0 +1,41 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_11 import ( + mo_velocity_advection_stencil_11, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil11(StencilTest): + PROGRAM = mo_velocity_advection_stencil_11 + OUTPUTS = ("z_w_con_c",) + + @staticmethod + def reference(mesh, w: np.array, **kwargs) -> dict: + z_w_con_c = w + return dict(z_w_con_c=z_w_con_c) + + @pytest.fixture + def input_data(self, mesh): + w = random_field(mesh, CellDim, KDim) + z_w_con_c = zero_field(mesh, CellDim, KDim) + return dict(w=w, z_w_con_c=z_w_con_c) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_12.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_12.py new file mode 100644 index 0000000000..b3dd9342de --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_12.py @@ -0,0 +1,39 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_12 import ( + mo_velocity_advection_stencil_12, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoVelocityAdvectionStencil12(StencilTest): + PROGRAM = mo_velocity_advection_stencil_12 + OUTPUTS = ("z_w_con_c",) + + @staticmethod + def reference(mesh, z_w_con_c: np.array, **kwargs) -> np.array: + z_w_con_c = np.zeros_like(z_w_con_c) + return dict(z_w_con_c=z_w_con_c) + + @pytest.fixture + def input_data(self, mesh): + z_w_con_c = random_field(mesh, CellDim, KDim) + + return dict( + z_w_con_c=z_w_con_c, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py new file mode 100644 index 0000000000..eb7fb3fc44 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py @@ -0,0 +1,41 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_13 import ( + mo_velocity_advection_stencil_13, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoVelocityAdvectionStencil13(StencilTest): + PROGRAM = mo_velocity_advection_stencil_13 + OUTPUTS = ("z_w_con_c",) + + @staticmethod + def reference(mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> np.array: + z_w_con_c = z_w_con_c - w_concorr_c + return dict(z_w_con_c=z_w_con_c) + + @pytest.fixture + def input_data(self, mesh): + z_w_con_c = random_field(mesh, CellDim, KDim) + w_concorr_c = random_field(mesh, CellDim, KDim) + + return dict( + w_concorr_c=w_concorr_c, + z_w_con_c=z_w_con_c, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py new file mode 100644 index 0000000000..f691e32ad2 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py @@ -0,0 +1,86 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_14 import ( + mo_velocity_advection_stencil_14, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + + +class TestMoVelocityAdvectionStencil14(StencilTest): + PROGRAM = mo_velocity_advection_stencil_14 + OUTPUTS = ("cfl_clipping", "pre_levelmask", "vcfl", "z_w_con_c") + + @staticmethod + def reference( + mesh, ddqz_z_half: np.array, z_w_con_c: np.array, cfl_w_limit, dtime, **kwargs + ) -> dict: + num_rows, num_cols = z_w_con_c.shape + cfl_clipping = np.where( + np.abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, + np.ones([num_rows, num_cols]), + np.zeros_like(z_w_con_c), + ) + num_rows, num_cols = cfl_clipping.shape + pre_levelmask = np.where( + cfl_clipping == 1.0, + np.ones([num_rows, num_cols]), + np.zeros_like(cfl_clipping), + ) + vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0) + z_w_con_c = np.where( + (cfl_clipping == 1.0) & (vcfl < -0.85), + -0.85 * ddqz_z_half / dtime, + z_w_con_c, + ) + z_w_con_c = np.where( + (cfl_clipping == 1.0) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c + ) + + return dict( + cfl_clipping=cfl_clipping, + pre_levelmask=pre_levelmask, + vcfl=vcfl, + z_w_con_c=z_w_con_c, + ) + + @pytest.fixture + def input_data(self, mesh): + ddqz_z_half = random_field(mesh, CellDim, KDim) + z_w_con_c = random_field(mesh, CellDim, KDim) + cfl_clipping = random_mask(mesh, CellDim, KDim, dtype=bool) + pre_levelmask = random_mask( + mesh, CellDim, KDim, dtype=bool + ) # TODO should be just a K field + vcfl = zero_field(mesh, CellDim, KDim) + cfl_w_limit = 5.0 + dtime = 9.0 + + return dict( + ddqz_z_half=ddqz_z_half, + z_w_con_c=z_w_con_c, + cfl_clipping=cfl_clipping, + pre_levelmask=pre_levelmask, + vcfl=vcfl, + cfl_w_limit=cfl_w_limit, + dtime=dtime, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_15.py new file mode 100644 index 0000000000..ef91adc33b --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_15.py @@ -0,0 +1,46 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_15 import ( + mo_velocity_advection_stencil_15, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestMoVelocityAdvectionStencil15(StencilTest): + PROGRAM = mo_velocity_advection_stencil_15 + OUTPUTS = ("z_w_con_c_full",) + + @staticmethod + def reference(mesh, z_w_con_c: np.array, **kwargs): + z_w_con_c_full = 0.5 * (z_w_con_c[:, :-1] + z_w_con_c[:, 1:]) + return dict(z_w_con_c_full=z_w_con_c_full) + + @pytest.fixture + def input_data(self, mesh): + z_w_con_c = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + + z_w_con_c_full = zero_field(mesh, CellDim, KDim) + + return dict( + z_w_con_c=z_w_con_c, + z_w_con_c_full=z_w_con_c_full, + ) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_16.py similarity index 86% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_16.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_16.py index a20167c385..4c402535bb 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_16.py @@ -13,13 +13,12 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_16 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_16 import ( mo_velocity_advection_stencil_16, ) -from icon4py.common.dimension import CellDim, KDim - -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_velocity_advection_stencil_16_numpy( diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_17.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_17.py new file mode 100644 index 0000000000..c265513244 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_17.py @@ -0,0 +1,46 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_17 import ( + mo_velocity_advection_stencil_17, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestMoVelocityAdvectionStencil17(StencilTest): + PROGRAM = mo_velocity_advection_stencil_17 + OUTPUTS = ("ddt_w_adv",) + + @staticmethod + def reference( + mesh, e_bln_c_s: np.array, z_v_grad_w: np.array, ddt_w_adv: np.array, **kwargs + ) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + ddt_w_adv = ddt_w_adv + np.sum(z_v_grad_w[mesh.c2e] * e_bln_c_s, axis=1) + return dict(ddt_w_adv=ddt_w_adv) + + @pytest.fixture + def input_data(self, mesh): + z_v_grad_w = random_field(mesh, EdgeDim, KDim) + e_bln_c_s = random_field(mesh, CellDim, C2EDim) + ddt_w_adv = random_field(mesh, CellDim, KDim) + + return dict( + e_bln_c_s=e_bln_c_s, + z_v_grad_w=z_v_grad_w, + ddt_w_adv=ddt_w_adv, + ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py new file mode 100644 index 0000000000..56d33d5454 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py @@ -0,0 +1,100 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_18 import ( + mo_velocity_advection_stencil_18, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, +) + + +class TestMoVelocityAdvectionStencil18(StencilTest): + PROGRAM = mo_velocity_advection_stencil_18 + OUTPUTS = ("ddt_w_adv",) + + @staticmethod + def reference( + mesh, + levelmask: np.array, + cfl_clipping: np.array, + owner_mask: np.array, + z_w_con_c: np.array, + ddqz_z_half: np.array, + area: np.array, + geofac_n2s: np.array, + w: np.array, + ddt_w_adv: np.array, + scalfac_exdiff: float, + cfl_w_limit: float, + dtime: float, + **kwargs, + ): + levelmask = np.expand_dims(levelmask, axis=0) + owner_mask = np.expand_dims(owner_mask, axis=-1) + area = np.expand_dims(area, axis=-1) + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + + difcoef = np.where( + (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), + scalfac_exdiff + * np.minimum( + 0.85 - cfl_w_limit * dtime, + np.abs(z_w_con_c) * dtime / ddqz_z_half - cfl_w_limit * dtime, + ), + 0, + ) + + ddt_w_adv = np.where( + (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), + ddt_w_adv + difcoef * area * np.sum(w[mesh.c2e2cO] * geofac_n2s, axis=1), + ddt_w_adv, + ) + + return dict(ddt_w_adv=ddt_w_adv) + + @pytest.fixture + def input_data(self, mesh): + levelmask = random_mask(mesh, KDim) + cfl_clipping = random_mask(mesh, CellDim, KDim) + owner_mask = random_mask(mesh, CellDim) + z_w_con_c = random_field(mesh, CellDim, KDim) + ddqz_z_half = random_field(mesh, CellDim, KDim) + area = random_field(mesh, CellDim) + geofac_n2s = random_field(mesh, CellDim, C2E2CODim) + w = random_field(mesh, CellDim, KDim) + ddt_w_adv = random_field(mesh, CellDim, KDim) + scalfac_exdiff = 10.0 + cfl_w_limit = 3.0 + dtime = 2.0 + + return dict( + levelmask=levelmask, + cfl_clipping=cfl_clipping, + owner_mask=owner_mask, + z_w_con_c=z_w_con_c, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + w=w, + ddt_w_adv=ddt_w_adv, + scalfac_exdiff=scalfac_exdiff, + cfl_w_limit=cfl_w_limit, + dtime=dtime, + ) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py similarity index 90% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py index d0c0a2839d..eaa6532beb 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py @@ -26,10 +26,10 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_19 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_19 import ( mo_velocity_advection_stencil_19, ) -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( CellDim, E2CDim, ECDim, @@ -37,9 +37,12 @@ KDim, VertexDim, ) - -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_velocity_advection_stencil_19_numpy( @@ -63,10 +66,7 @@ def mo_velocity_advection_stencil_19_numpy( ddt_vn_adv = -( (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + ( - -coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] - + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1] - ) + + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + vt * (f_e + 0.5 * np.sum(zeta[e2v], axis=1)) + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) * (vn_ie[:, :-1] - vn_ie[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py similarity index 94% rename from atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py rename to model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py index 101628edba..5a37883c2c 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py @@ -13,10 +13,10 @@ import numpy as np -from icon4py.atm_dyn_iconam.mo_velocity_advection_stencil_20 import ( +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_20 import ( mo_velocity_advection_stencil_20, ) -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( CellDim, E2C2EODim, E2CDim, @@ -24,9 +24,8 @@ KDim, VertexDim, ) - -from .test_utils.helpers import random_field, random_mask -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, random_mask +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_velocity_advection_stencil_20_numpy( diff --git a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py new file mode 100644 index 0000000000..bae0e5f16c --- /dev/null +++ b/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -0,0 +1,59 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( + temporary_field_for_grid_point_cold_pools_enhancement, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestTemporaryFieldForGridPointColdPoolsEnhancement(StencilTest): + PROGRAM = temporary_field_for_grid_point_cold_pools_enhancement + OUTPUTS = ("enh_diffu_3d",) + + @staticmethod + def reference( + mesh, theta_v: np.array, theta_ref_mc: np.array, thresh_tdiff, **kwargs + ) -> np.array: + tdiff = theta_v - np.sum(theta_v[mesh.c2e2c], axis=1) / 3 + trefdiff = theta_ref_mc - np.sum(theta_ref_mc[mesh.c2e2c], axis=1) / 3 + + enh_diffu_3d = np.where( + ((tdiff - trefdiff) < thresh_tdiff) & (trefdiff < 0), + (thresh_tdiff - tdiff + trefdiff) * 5e-4, + -1.7976931348623157e308, + ) + + return dict(enh_diffu_3d=enh_diffu_3d) + + @pytest.fixture + def input_data(self, mesh): + theta_v = random_field(mesh, CellDim, KDim) + theta_ref_mc = random_field(mesh, CellDim, KDim) + enh_diffu_3d = zero_field(mesh, CellDim, KDim) + thresh_tdiff = 5.0 + + return dict( + theta_v=theta_v, + theta_ref_mc=theta_ref_mc, + enh_diffu_3d=enh_diffu_3d, + thresh_tdiff=thresh_tdiff, + ) diff --git a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py new file mode 100644 index 0000000000..268db90575 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -0,0 +1,73 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( + temporary_fields_for_turbulence_diagnostics, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + zero_field, +) + + +class TestTemporaryFieldsForTurbulenceDiagnostics(StencilTest): + PROGRAM = temporary_fields_for_turbulence_diagnostics + OUTPUTS = ("div", "kh_c") + + @staticmethod + def reference( + mesh, + kh_smag_ec: np.array, + vn: np.array, + e_bln_c_s: np.array, + geofac_div: np.array, + diff_multfac_smag: np.array, + **kwargs, + ) -> dict: + geofac_div = np.expand_dims(geofac_div, axis=-1) + vn_geofac = vn[mesh.c2e] * geofac_div + div = np.sum(vn_geofac, axis=1) + + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + diff_multfac_smag = np.expand_dims(diff_multfac_smag, axis=0) + mul = kh_smag_ec[mesh.c2e] * e_bln_c_s + summed = np.sum(mul, axis=1) + kh_c = summed / diff_multfac_smag + + return dict(div=div, kh_c=kh_c) + + @pytest.fixture + def input_data(self, mesh): + vn = random_field(mesh, EdgeDim, KDim) + geofac_div = random_field(mesh, CellDim, C2EDim) + kh_smag_ec = random_field(mesh, EdgeDim, KDim) + e_bln_c_s = random_field(mesh, CellDim, C2EDim) + diff_multfac_smag = random_field(mesh, KDim) + + kh_c = zero_field(mesh, CellDim, KDim) + div = zero_field(mesh, CellDim, KDim) + + return dict( + kh_smag_ec=kh_smag_ec, + vn=vn, + e_bln_c_s=e_bln_c_s, + geofac_div=geofac_div, + diff_multfac_smag=diff_multfac_smag, + kh_c=kh_c, + div=div, + ) diff --git a/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 86% rename from atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 28ee4fff36..123a8ab301 100644 --- a/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,18 +15,17 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim - -from .test_utils.helpers import ( +from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( flatten_first_two_dims, random_field, random_mask, zero_field, ) -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def mo_nh_diffusion_stencil_15_numpy( @@ -57,15 +56,12 @@ def mo_nh_diffusion_stencil_15_numpy( ] sum_over = np.sum( - geofac_n2s_nbh - * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), axis=1, ) geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - return np.where( - mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp - ) + return np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) def test_mo_nh_diffusion_stencil_15(): @@ -106,11 +102,6 @@ def test_mo_nh_diffusion_stencil_15(): np.asarray(z_temp), ) - hstart = 0 - hend = mesh.n_cells - kstart = 0 - kend = mesh.k_level - truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, zd_vertoffset_new, @@ -120,10 +111,6 @@ def test_mo_nh_diffusion_stencil_15(): vcoef_new, theta_v, z_temp, - hstart, - hend, - kstart, - kend, offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), diff --git a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py b/model/atmosphere/dycore/tests/test_update_theta_and_exner.py new file mode 100644 index 0000000000..af6d6062fe --- /dev/null +++ b/model/atmosphere/dycore/tests/test_update_theta_and_exner.py @@ -0,0 +1,58 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.update_theta_and_exner import ( + update_theta_and_exner, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestUpdateThetaAndExner(StencilTest): + PROGRAM = update_theta_and_exner + OUTPUTS = ("theta_v", "exner") + + @staticmethod + def reference( + mesh, + z_temp: np.array, + area: np.array, + theta_v: np.array, + exner: np.array, + rd_o_cvd, + **kwargs, + ) -> tuple[np.array]: + area = np.expand_dims(area, axis=0) + z_theta = theta_v + theta_v = theta_v + (np.expand_dims(area, axis=-1) * z_temp) + exner = exner * (1.0 + rd_o_cvd * (theta_v / z_theta - 1.0)) + return dict(theta_v=theta_v, exner=exner) + + @pytest.fixture + def input_data(self, mesh): + z_temp = random_field(mesh, CellDim, KDim) + area = random_field(mesh, CellDim) + theta_v = random_field(mesh, CellDim, KDim) + exner = random_field(mesh, CellDim, KDim) + rd_o_cvd = 5.0 + + return dict( + z_temp=z_temp, + area=area, + theta_v=theta_v, + exner=exner, + rd_o_cvd=rd_o_cvd, + ) diff --git a/model/common/.bumpversion.cfg b/model/common/.bumpversion.cfg new file mode 100644 index 0000000000..777683b350 --- /dev/null +++ b/model/common/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/common/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/common/README.md b/model/common/README.md new file mode 100644 index 0000000000..aa6a42598a --- /dev/null +++ b/model/common/README.md @@ -0,0 +1,9 @@ +# icon4py-common + +## Description + +Utilities shared by several ICON4Py components. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml new file mode 100644 index 0000000000..8fae7b1c1e --- /dev/null +++ b/model/common/pyproject.toml @@ -0,0 +1,117 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel>=0.40.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "icon4py-common" +description = "Shared code for the icon4py model." +readme = "README.md" +requires-python = ">=3.10" +license = {file = "LICENSE"} +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", +] +dynamic = ['version'] + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +parallel = true +branch = true +source_pkgs = ['common'] + +[tool.isort] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +multi_line_output = 3 +use_parentheses = true +include_trailing_comma = true +force_grid_wrap = 0 + +[tool.mypy] +install_types = true +non_interactive = true +exclude = [ + '^tests/*.py', +] +disallow_incomplete_defs = true +disallow_untyped_defs = true +ignore_missing_imports = false +implicit_reexport = true +warn_unused_configs = true +warn_unused_ignores = true +warn_redundant_casts = true +show_column_numbers = true +show_error_codes = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.common.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.common' = ['py.typed'] diff --git a/model/common/requirements-dev.txt b/model/common/requirements-dev.txt new file mode 100644 index 0000000000..7cad1ce7ca --- /dev/null +++ b/model/common/requirements-dev.txt @@ -0,0 +1,2 @@ +-r ../../base-requirements-dev.txt +-e . diff --git a/model/common/requirements.txt b/model/common/requirements.txt new file mode 100644 index 0000000000..58cc42cc02 --- /dev/null +++ b/model/common/requirements.txt @@ -0,0 +1,2 @@ +-r ../../base-requirements.txt +. diff --git a/atm_dyn_iconam/setup.py b/model/common/src/icon4py/model/common/__init__.py similarity index 55% rename from atm_dyn_iconam/setup.py rename to model/common/src/icon4py/model/common/__init__.py index 9c9f7b81c8..dab7089554 100644 --- a/atm_dyn_iconam/setup.py +++ b/model/common/src/icon4py/model/common/__init__.py @@ -10,9 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final -from setuptools import setup +from packaging import version as pkg_version -if __name__ == "__main__": - setup() +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/common/src/icon4py/common/dimension.py b/model/common/src/icon4py/model/common/dimension.py similarity index 100% rename from common/src/icon4py/common/dimension.py rename to model/common/src/icon4py/model/common/dimension.py diff --git a/tools/src/icon4pytools/icon4pygen/py.typed b/model/common/src/icon4py/model/common/py.typed similarity index 100% rename from tools/src/icon4pytools/icon4pygen/py.typed rename to model/common/src/icon4py/model/common/py.typed diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/__init__.py b/model/common/src/icon4py/model/common/test_utils/__init__.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/__init__.py rename to model/common/src/icon4py/model/common/test_utils/__init__.py diff --git a/atm_dyn_iconam/tests/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py similarity index 58% rename from atm_dyn_iconam/tests/test_utils/helpers.py rename to model/common/src/icon4py/model/common/test_utils/helpers.py index f85bbf8f58..eee6ac3cfc 100644 --- a/atm_dyn_iconam/tests/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -11,13 +11,21 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import ClassVar from typing import Optional import numpy as np import numpy.typing as npt +import pytest from gt4py.next import common as gt_common +from gt4py.next.ffront.decorator import Program from gt4py.next.iterator import embedded as it_embedded +try: + import pytest_benchmark +except ModuleNotFoundError: + pytest_benchmark = None + from .simple_mesh import SimpleMesh @@ -58,9 +66,7 @@ def random_field( extend: Optional[dict[gt_common.Dimension, int]] = None, ) -> it_embedded.MutableLocatedField: return it_embedded.np_as_located_field(*dims)( - np.random.default_rng().uniform( - low=low, high=high, size=_shape(mesh, *dims, extend=extend) - ) + np.random.default_rng().uniform(low=low, high=high, size=_shape(mesh, *dims, extend=extend)) ) @@ -104,3 +110,66 @@ def flatten_first_two_dims( new_shape = flattened_shape + old_shape[2:] newarray = np.asarray(field).reshape(new_shape) return it_embedded.np_as_located_field(*dims)(newarray) + + +def _test_validation(self, mesh, backend, input_data): + reference_outputs = self.reference(mesh, **{k: np.array(v) for k, v in input_data.items()}) + self.PROGRAM.with_backend(backend)( + **input_data, + offset_provider=mesh.get_offset_provider(), + ) + for name in self.OUTPUTS: + assert np.allclose( + input_data[name], reference_outputs[name] + ), f"Validation failed for '{name}'" + + +if pytest_benchmark: + + def _bench_execution(self, pytestconfig, mesh, backend, input_data, benchmark): + if pytestconfig.getoption( + "--benchmark-disable" + ): # skipping as otherwise program calls are duplicated in tests. + pytest.skip("Test skipped due to 'benchmark-disable' option.") + else: + benchmark( + self.PROGRAM.with_backend(backend), + **input_data, + offset_provider=mesh.get_offset_provider(), + ) + +else: + + def _bench_execution(self, pytestconfig): + pytest.skip("Test skipped as `pytest-benchmark` is not installed.") + + +class StencilTest: + """ + Base class to be used for testing stencils. + + Example (pseudo-code): + + >>> class TestMultiplyByTwo(StencilTest): # doctest: +SKIP + ... PROGRAM = multiply_by_two # noqa: F821 + ... OUTPUTS = ("some_output",) + ... + ... @pytest.fixture + ... def input_data(self): + ... return {"some_input": ..., "some_output": ...} + ... + ... @staticmethod + ... def reference(some_input, **kwargs): + ... return dict(some_output=np.asarray(some_input)*2) + """ + + PROGRAM: ClassVar[Program] + OUTPUTS: ClassVar[tuple[str, ...]] + + def __init_subclass__(cls, **kwargs): + # Add two methods for verification and benchmarking. In order to have names that + # reflect the name of the test we do this dynamically here instead of using regular + # inheritance. + super().__init_subclass__(**kwargs) + setattr(cls, f"test_{cls.__name__}", _test_validation) + setattr(cls, f"bench_{cls.__name__}", _bench_execution) diff --git a/atm_dyn_iconam/tests/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py similarity index 94% rename from atm_dyn_iconam/tests/test_utils/simple_mesh.py rename to model/common/src/icon4py/model/common/test_utils/simple_mesh.py index 5b66aeed19..2cff62c6eb 100644 --- a/atm_dyn_iconam/tests/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -16,7 +16,7 @@ import numpy as np from gt4py.next.iterator.embedded import NeighborTableOffsetProvider -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, C2EDim, @@ -425,3 +425,18 @@ def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) + + def get_offset_provider(self): + return { + "C2E": self.get_c2e_offset_provider(), + "C2E2CO": self.get_c2e2cO_offset_provider(), + "C2E2C": self.get_c2e2c_offset_provider(), + "E2C2EO": self.get_e2c2eO_offset_provider(), + "E2C2E": self.get_e2c2e_offset_provider(), + "V2C": self.get_v2c_offset_provider(), + "V2E": self.get_v2e_offset_provider(), + "E2C": self.get_e2c_offset_provider(), + "E2V": self.get_e2v_offset_provider(), + "E2C2V": self.get_e2c2v_offset_provider(), + "Koff": KDim, + } diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt new file mode 100644 index 0000000000..1b7e16d5a1 --- /dev/null +++ b/model/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../base-requirements-dev.txt +-e ./atmosphere/dycore +-e ./common diff --git a/model/requirements.txt b/model/requirements.txt new file mode 100644 index 0000000000..2045fe9058 --- /dev/null +++ b/model/requirements.txt @@ -0,0 +1,3 @@ +-r ../base-requirements.txt +./atmosphere/dycore +./common diff --git a/model/tox.ini b/model/tox.ini new file mode 100644 index 0000000000..27fb18d610 --- /dev/null +++ b/model/tox.ini @@ -0,0 +1,33 @@ +# Tox configuration file +# Reference documentation: https://tox.readthedocs.org/ + +[tox] +envlist = + py{310} + dev +skipsdist = true + +[testenv] +passenv = + PIP_USER + PYTHONUSERBASE +deps = + -r {toxinidir}/requirements-dev.txt +commands = + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src + pytest -v -s -n auto --cov --cov-append --benchmark-disable +commands_post = + rm -rf tests/_reports/coverage_html + -coverage html + -coverage report +allowlist_externals = + /bin/bash + rm + +[testenv:dev] +basepython = python3.10 +setenv = + PIP_SRC = _external_src +skip_install = true +commands = +commands_post = diff --git a/pytest.ini b/pytest.ini index d9da53905e..84a9d34b4e 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,3 @@ [pytest] norecursedirs = _external_src +python_functions = test_* bench_* diff --git a/requirements-dev.txt b/requirements-dev.txt index ae585c1cef..809c77a57b 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,9 +1,9 @@ -r base-requirements-dev.txt -# Concretized dependencies in local repo --e ./common --e ./atm_dyn_iconam --e ./advection +# icon4py model +-e ./model/atmosphere/dycore +-e ./model/atmosphere/advection +-e ./model/common # icon4pytools -e ./tools diff --git a/requirements.txt b/requirements.txt index 22a84af692..ca7931901c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,9 @@ -r base-requirements.txt -# Concretized dependencies in local repo -./common -./atm_dyn_iconam -./advection +# icon4py model +./model/atmosphere/dycore +./model/atmosphere/advection +./model/common # icon4pytools ./tools diff --git a/spack/gt4py-stable/spack.yaml b/spack/gt4py-stable/spack.yaml index b1b847cc97..d4a93e8bab 100644 --- a/spack/gt4py-stable/spack.yaml +++ b/spack/gt4py-stable/spack.yaml @@ -1,11 +1,10 @@ spack: specs: - - py-icon4py@main%gcc@9.3.0 - - py-gt4py@1.1.1%gcc@9.3.0 + - py-icon4py@main%gcc@9.3.0 ^py-gt4py@1.1.2%gcc@9.3.0 view: false concretizer: unify: true develop: py-icon4py: - spec: py-icon4py@main%gcc@9.3.0 + spec: py-icon4py@main%gcc@9.3.0 ^py-gt4py@1.1.2%gcc@9.3.0 path: ../../ diff --git a/tools/.bumpversion.cfg b/tools/.bumpversion.cfg index c04dfc55bd..62fb12e51e 100644 --- a/tools/.bumpversion.cfg +++ b/tools/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.0.4 +current_version = 0.0.6 parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? serialize = {major}.{minor}.{patch} diff --git a/tools/.license_header.txt b/tools/.license_header.txt new file mode 100644 index 0000000000..8e88814589 --- /dev/null +++ b/tools/.license_header.txt @@ -0,0 +1,12 @@ +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 diff --git a/tools/.pre-commit-config.yaml b/tools/.pre-commit-config.yaml index 533cbae3f9..851ca38e5b 100644 --- a/tools/.pre-commit-config.yaml +++ b/tools/.pre-commit-config.yaml @@ -5,7 +5,7 @@ default_language_version: # Remove frozen version once we migrated away from tsa node: 17.9.1 minimum_pre_commit_version: 2.20.0 -exclude: "advection/.*", "atm_dyn_iconam/.*" +exclude: "model/.*" repos: - repo: meta @@ -52,7 +52,7 @@ repos: --comment-style, "|#|", --license-filepath, - .license_header.txt, + tools/.license_header.txt, --fuzzy-match-generates-todo, ] @@ -110,6 +110,7 @@ repos: - flake8-eradicate - flake8-mutable - pygments + args: [--config=tools/.flake8, tools/src/icon4pytools] - repo: local hooks: diff --git a/tools/pyproject.toml b/tools/pyproject.toml index b8a6e36535..e8b4f6491e 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -27,8 +27,8 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "icon4py_advection", - "icon4py_atm_dyn_iconam", + "icon4py-atmosphere-dycore", + "icon4py-atmosphere-advection", "gt4py>=1.0.1", "icon4py_common", "tabulate>=0.8.9", diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index 7b1ad57ba8..4dc72f7d88 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,5 +1,5 @@ -r ../base-requirements-dev.txt --e ../common --e ../advection --e ../atm_dyn_iconam +-e ../model/atmosphere/dycore +-e ../model/atmosphere/advection +-e ../model/common -e . diff --git a/tools/requirements.txt b/tools/requirements.txt index e9c0e49500..dfe3e8e4e7 100644 --- a/tools/requirements.txt +++ b/tools/requirements.txt @@ -1,5 +1,5 @@ -r ../base-requirements.txt -../common -../advection -../atm_dyn_iconam +../model/atmosphere/dycore +../model/atmosphere/advection +../model/common . diff --git a/tools/src/icon4pytools/__init__.py b/tools/src/icon4pytools/__init__.py index b8a1e792ed..ae71b6e01d 100644 --- a/tools/src/icon4pytools/__init__.py +++ b/tools/src/icon4pytools/__init__.py @@ -32,5 +32,5 @@ __license__: Final = "GPL-3.0-or-later" -__version__: Final = "0.0.4" +__version__: Final = "0.0.6" __version_info__: Final = pkg_version.parse(__version__) diff --git a/tools/src/icon4pytools/common/__init__.py b/tools/src/icon4pytools/common/__init__.py index 15dfdb0098..0da2c55dd5 100644 --- a/tools/src/icon4pytools/common/__init__.py +++ b/tools/src/icon4pytools/common/__init__.py @@ -10,3 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later + +ICON4PY_MODEL_QUALIFIED_NAME = "icon4py.model" + +__all__ = ["ICON4PY_MODEL_QUALIFIED_NAME"] diff --git a/tools/src/icon4pytools/f2ser/parse.py b/tools/src/icon4pytools/f2ser/parse.py index c4addb3870..a0f16a5838 100644 --- a/tools/src/icon4pytools/f2ser/parse.py +++ b/tools/src/icon4pytools/f2ser/parse.py @@ -65,9 +65,38 @@ def __init__(self, granule: Path, dependencies: Optional[list[Path]] = None) -> def __call__(self) -> ParsedGranule: """Parse the granule and return the parsed data.""" subroutines = self.parse_subroutines() - last_import_ln = self.find_last_fortran_use_statement() + last_import_ln = self._find_last_fortran_use_statement() return ParsedGranule(subroutines=subroutines, last_import_ln=last_import_ln) + def _find_last_fortran_use_statement(self) -> Optional[int]: + """Find the line number of the last Fortran USE statement in the code. + + Returns: + int: the line number of the last USE statement, or None if no USE statement is found. + """ + # Reverse the order of the lines so we can search from the end + code = self._read_code_from_file() + code_lines = code.splitlines() + code_lines.reverse() + + # Look for the last USE statement + use_ln = None + for i, line in enumerate(code_lines): + if line.strip().lower().startswith("use"): + use_ln = len(code_lines) - i + if i > 0 and code_lines[i - 1].strip().lower() == "#endif": + # If the USE statement is preceded by an #endif statement, return the line number after the #endif statement + return use_ln + 1 + else: + return use_ln + return None + + def _read_code_from_file(self) -> str: + """Read the content of the granule and returns it as a string.""" + with open(self.granule_path) as f: + code = f.read() + return code + def parse_subroutines(self): subroutines = self._extract_subroutines(crack(self.granule_path)) variables_grouped_by_intent = { @@ -245,31 +274,12 @@ def _update_with_codegen_lines(self, parsed_types: dict) -> dict: with_lines = deepcopy(parsed_types) for subroutine in with_lines: for intent in with_lines[subroutine]: - with_lines[subroutine][intent]["codegen_ctx"] = self.get_subroutine_lines( + with_lines[subroutine][intent]["codegen_ctx"] = self._get_subroutine_lines( subroutine ) return with_lines - def find_last_fortran_use_statement(self): - with open(self.granule_path) as f: - file_contents = f.readlines() - - # Reverse the order of the lines so we can search from the end - file_contents.reverse() - - # Look for the last USE statement - use_ln = None - for i, line in enumerate(file_contents): - if line.strip().lower().startswith("use"): - use_ln = len(file_contents) - i - if i > 0 and file_contents[i - 1].strip().lower() == "#endif": - # If the USE statement is preceded by an #endif statement, return the line number after the #endif statement - return use_ln + 1 - else: - return use_ln - return None - - def get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: + def _get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: """Return CodegenContext object containing line numbers of the last declaration statement and the code before the end of the given subroutine. Args: @@ -278,10 +288,39 @@ def get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: Returns: CodegenContext: Object containing the line number of the last declaration statement and the line number of the last line of the code before the end of the given subroutine. """ - with open(self.granule_path) as f: - code = f.read() + code = self._read_code_from_file() + + start_subroutine_ln, end_subroutine_ln = self._find_subroutine_lines(code, subroutine_name) + + variable_declaration_ln = self._find_variable_declarations( + code, start_subroutine_ln, end_subroutine_ln + ) - # Find the line number where the subroutine is defined + if not variable_declaration_ln: + raise ParsingError(f"No variable declarations found in {self.granule_path}") + + ( + first_declaration_ln, + last_declaration_ln, + ) = self._get_variable_declaration_bounds(variable_declaration_ln, start_subroutine_ln) + + pre_end_subroutine_ln = ( + end_subroutine_ln - 1 + ) # we want to generate the code before the end of the subroutine + + return CodegenContext(first_declaration_ln, last_declaration_ln, pre_end_subroutine_ln) + + @staticmethod + def _find_subroutine_lines(code: str, subroutine_name: str) -> tuple[int]: + """Find line numbers of a subroutine within a code block. + + Args: + code (str): The code block to search for the subroutine. + subroutine_name (str): Name of the subroutine to find. + + Returns: + tuple: Line numbers of the start and end of the subroutine. + """ start_subroutine_pattern = r"SUBROUTINE\s+" + subroutine_name + r"\s*\(" end_subroutine_pattern = r"END\s+SUBROUTINE\s+" + subroutine_name + r"\s*" start_match = re.search(start_subroutine_pattern, code) @@ -290,21 +329,58 @@ def get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: return None start_subroutine_ln = code[: start_match.start()].count("\n") + 1 end_subroutine_ln = code[: end_match.start()].count("\n") + 1 + return start_subroutine_ln, end_subroutine_ln - # Find the last intent statement line number in the subroutine - declaration_pattern = r".*::\s*(\w+\b)" - declaration_pattern_lines = [ - i - for i, line in enumerate(code.splitlines()[start_subroutine_ln:end_subroutine_ln]) - if re.search(declaration_pattern, line) - ] - if not declaration_pattern_lines: - raise ParsingError(f"No declarations found in {self.granule_path}") - last_declaration_ln = declaration_pattern_lines[-1] + start_subroutine_ln + 1 - first_declaration_ln = declaration_pattern_lines[0] + start_subroutine_ln + @staticmethod + def _find_variable_declarations( + code: str, start_subroutine_ln: int, end_subroutine_ln: int + ) -> list: + """Find line numbers of variable declarations within a code block. - pre_end_subroutine_ln = ( - end_subroutine_ln - 1 - ) # we want to generate the code before the end of the subroutine + Args: + code (str): The code block to search for variable declarations. + start_subroutine_ln (int): Starting line number of the subroutine. + end_subroutine_ln (int): Ending line number of the subroutine. - return CodegenContext(first_declaration_ln, last_declaration_ln, pre_end_subroutine_ln) + Returns: + list: Line numbers of variable declaration lines. + + This method identifies single-line and multiline variable declarations within + the specified code block, delimited by the start and end line numbers of the + subroutine. Multiline declarations are detected by the presence of an ampersand + character ('&') at the end of a line. + """ + declaration_pattern = r".*::\s*(\w+\b)|.*::.*(\&)" + is_multiline_declaration = False + declaration_pattern_lines = [] + + for i, line in enumerate(code.splitlines()[start_subroutine_ln:end_subroutine_ln]): + if not is_multiline_declaration: + if re.search(declaration_pattern, line): + declaration_pattern_lines.append(i) + if line.find("&") != -1: + is_multiline_declaration = True + else: + if is_multiline_declaration: + declaration_pattern_lines.append(i) + if line.find("&") == -1: + is_multiline_declaration = False + + return declaration_pattern_lines + + @staticmethod + def _get_variable_declaration_bounds( + declaration_pattern_lines: list, start_subroutine_ln: int + ) -> tuple: + """Return the line numbers of the bounds for a variable declaration block. + + Args: + declaration_pattern_lines (list): List of line numbers representing the relative positions of lines within the declaration block. + start_subroutine_ln (int): Line number indicating the starting line of the subroutine. + + Returns: + tuple: Line number of the first declaration line, line number following the last declaration line. + """ + first_declaration_ln = declaration_pattern_lines[0] + start_subroutine_ln + last_declaration_ln = declaration_pattern_lines[-1] + start_subroutine_ln + 1 + return first_declaration_ln, last_declaration_ln diff --git a/tools/src/icon4pytools/icon4pygen/backend.py b/tools/src/icon4pytools/icon4pygen/backend.py index 80b5a10e9a..01b917afd1 100644 --- a/tools/src/icon4pytools/icon4pygen/backend.py +++ b/tools/src/icon4pytools/icon4pygen/backend.py @@ -15,7 +15,7 @@ from gt4py.next.iterator import ir as itir from gt4py.next.program_processors.codegens.gtfn.gtfn_backend import generate -from icon4py.common.dimension import Koff +from icon4py.model.common.dimension import Koff from icon4pytools.icon4pygen.bindings.utils import write_string from icon4pytools.icon4pygen.exceptions import MultipleFieldOperatorException diff --git a/tools/src/icon4pytools/icon4pygen/bindings/entities.py b/tools/src/icon4pytools/icon4pygen/bindings/entities.py index a9297d9c5f..cad08c7c0d 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/entities.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/entities.py @@ -113,7 +113,6 @@ def is_compound(self) -> bool: def is_integral(self) -> bool: scalar_types = [ - ts.ScalarKind.INT, ts.ScalarKind.INT32, ts.ScalarKind.INT64, ] diff --git a/tools/src/icon4pytools/icon4pygen/bindings/locations.py b/tools/src/icon4pytools/icon4pygen/bindings/locations.py index 5a35054e75..37e3ef7926 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/locations.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/locations.py @@ -15,9 +15,10 @@ from typing import Iterator from gt4py.next.ffront.fbuiltins import Dimension +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim + from icon4pytools.icon4pygen.bindings.codegen.render.location import LocationRenderer -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim class BasicLocation: diff --git a/tools/src/icon4pytools/icon4pygen/cli.py b/tools/src/icon4pytools/icon4pygen/cli.py index 66180adf44..ef886a1a41 100644 --- a/tools/src/icon4pytools/icon4pygen/cli.py +++ b/tools/src/icon4pytools/icon4pygen/cli.py @@ -21,10 +21,11 @@ class ModuleType(click.ParamType): + dycore_import_path = "icon4py.model.atmosphere.dycore" names = [ - "icon4pytools.atm_dyn_iconam.mo_nh_diffusion_stencil_", - "icon4pytools.atm_dyn_iconam.mo_solve_nonhydro_stencil_", - "icon4pytools.atm_dyn_iconam.mo_velocity_advection_stencil_", + f"{dycore_import_path}.mo_nh_diffusion_stencil_", + f"{dycore_import_path}.mo_solve_nonhydro_stencil_", + f"{dycore_import_path}.mo_velocity_advection_stencil_", ] def shell_complete(self, ctx, param, incomplete): diff --git a/tools/src/icon4pytools/icon4pygen/icochainsize.py b/tools/src/icon4pytools/icon4pygen/icochainsize.py index 66b220a1cc..183b0a89fe 100644 --- a/tools/src/icon4pytools/icon4pygen/icochainsize.py +++ b/tools/src/icon4pytools/icon4pygen/icochainsize.py @@ -111,8 +111,7 @@ from typing import List, TypeAlias from gt4py.next.common import Dimension - -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim @dataclass diff --git a/tools/src/icon4pytools/icon4pygen/metadata.py b/tools/src/icon4pytools/icon4pygen/metadata.py index 340f811993..02edeba212 100644 --- a/tools/src/icon4pytools/icon4pygen/metadata.py +++ b/tools/src/icon4pytools/icon4pygen/metadata.py @@ -25,13 +25,14 @@ from gt4py.next.iterator import ir as itir from gt4py.next.iterator.runtime import FendefDispatcher from gt4py.next.type_system import type_specifications as ts +from icon4py.model.common.dimension import CellDim, EdgeDim, Koff, VertexDim + from icon4pytools.icon4pygen.exceptions import ( InvalidConnectivityException, MultipleFieldOperatorException, ) from icon4pytools.icon4pygen.icochainsize import IcoChainSize -from icon4py.common.dimension import CellDim, EdgeDim, Koff, VertexDim @dataclass(frozen=True) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 1884190504..4b4605c7cb 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -139,7 +139,6 @@ class StartCreateDataFactory(DataFactoryBase): dtype: Type[StartCreateData] = StartCreateData def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData]: - deserialised = [] extracted = extract_directive(parsed["directives"], self.directive_cls) @@ -147,7 +146,6 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData]: return UnusedDirective for i, directive in enumerate(extracted): - named_args = parsed["content"]["StartCreate"][i] extra_fields = None diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 755194875e..1b98413f77 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -16,6 +16,8 @@ from typing import Any from gt4py.next.ffront.decorator import Program + +from icon4pytools.common import ICON4PY_MODEL_QUALIFIED_NAME from icon4pytools.common.logger import setup_logger from icon4pytools.icon4pygen.metadata import get_stencil_info from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface @@ -27,7 +29,7 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atm_dyn_iconam", "advection"] + _STENCIL_PACKAGES = ["atmosphere.dycore"] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed @@ -44,7 +46,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: field_info = gt4py_fields[f.variable] except KeyError: raise IncompatibleFieldError( - f"Used field variable name that is incompatible with the expected field names defined in {s.name} in icon4pytools." + f"Used field variable name that is incompatible with the expected field names defined in {s.name} in icon4py." ) f.out = field_info.out f.inp = field_info.inp @@ -55,21 +57,20 @@ def _collect_icon4py_stencil(self, stencil_name: str) -> Program: err_counter = 0 for pkg in self._STENCIL_PACKAGES: try: - module_name = f"icon4py.{pkg}.{stencil_name}" + module_name = f"{ICON4PY_MODEL_QUALIFIED_NAME}.{pkg}.{stencil_name}" module = importlib.import_module(module_name) except ModuleNotFoundError: err_counter += 1 if err_counter == len(self._STENCIL_PACKAGES): - error_msg = f"Did not find module: {stencil_name} in icon4pytools." - raise UnknownStencilError(error_msg) + raise UnknownStencilError(f"Did not find module: {stencil_name}") module_members = getmembers(module) found_stencil = [elt for elt in module_members if elt[0] == stencil_name] if len(found_stencil) == 0: raise UnknownStencilError( - f"Did not find module member: {stencil_name} in module: {module.__name__} in icon4pytools." + f"Did not find module member: {stencil_name} in module: {module.__name__}" ) return found_stencil[0][1] diff --git a/tools/src/icon4pytools/liskov/py.typed b/tools/src/icon4pytools/liskov/py.typed deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tools/tests/f2ser/fortran_samples/multiline_example.f90 b/tools/tests/f2ser/fortran_samples/multiline_example.f90 new file mode 100644 index 0000000000..c3d403310f --- /dev/null +++ b/tools/tests/f2ser/fortran_samples/multiline_example.f90 @@ -0,0 +1,423 @@ +MODULE mo_graupel_granule + + USE, INTRINSIC :: iso_fortran_env, ONLY: wp => real64 + USE gscp_graupel, ONLY: graupel + + IMPLICIT NONE + + PRIVATE :: graupel_parameters, params, mma,mmb + PUBLIC :: graupel_init, graupel_run + + TYPE graupel_parameters + ! gscp_data + INTEGER :: iautocon + INTEGER :: isnow_n0temp + REAL(wp) :: ccsrim + REAL(wp) :: ccsagg + REAL(wp) :: ccsdep + REAL(wp) :: ccsvel + REAL(wp) :: ccsvxp + REAL(wp) :: ccslam + REAL(wp) :: ccslxp + REAL(wp) :: ccsaxp + REAL(wp) :: ccsdxp + REAL(wp) :: ccshi1 + REAL(wp) :: ccdvtp + REAL(wp) :: ccidep + REAL(wp) :: ccswxp + REAL(wp) :: zconst + REAL(wp) :: zcev + REAL(wp) :: zbev + REAL(wp) :: zcevxp + REAL(wp) :: zbevxp + REAL(wp) :: zvzxp + REAL(wp) :: zvz0r + REAL(wp) :: v0snow + REAL(wp) :: x13o8 + REAL(wp) :: x1o2 + REAL(wp) :: x27o16 + REAL(wp) :: x3o4 + REAL(wp) :: x7o4 + REAL(wp) :: x7o8 + REAL(wp) :: zbvi + REAL(wp) :: zcac + REAL(wp) :: zccau + REAL(wp) :: zciau + REAL(wp) :: zcicri + REAL(wp) :: zcrcri + REAL(wp) :: zcrfrz + REAL(wp) :: zcrfrz1 + REAL(wp) :: zcrfrz2 + REAL(wp) :: zeps + REAL(wp) :: zkcac + REAL(wp) :: zkphi1 + REAL(wp) :: zkphi2 + REAL(wp) :: zkphi3 + REAL(wp) :: zmi0 + REAL(wp) :: zmimax + REAL(wp) :: zmsmin + REAL(wp) :: zn0s0 + REAL(wp) :: zn0s1 + REAL(wp) :: zn0s2 + REAL(wp) :: znimax_thom + REAL(wp) :: zqmin + REAL(wp) :: zrho0 + REAL(wp) :: zthet + REAL(wp) :: zthn + REAL(wp) :: ztmix + REAL(wp) :: ztrfrz + REAL(wp) :: zvz0i + REAL(wp) :: icesedi_exp + REAL(wp) :: zams + REAL(wp) :: dist_cldtop_ref + REAL(wp) :: reduce_dep_ref + REAL(wp) :: tmin_iceautoconv + REAL(wp) :: zceff_fac + REAL(wp) :: zceff_min + REAL(wp) :: v_sedi_rain_min + REAL(wp) :: v_sedi_snow_min + REAL(wp) :: v_sedi_graupel_min + + ! mo_physical constants + REAL(wp) :: r_v + REAL(wp) :: lh_v + REAL(wp) :: lh_s + REAL(wp) :: cpdr + REAL(wp) :: cvdr + REAL(wp) :: b3 + REAL(wp) :: t0 + END TYPE graupel_parameters + + + + REAL(wp):: mma(10),mmb(10) + TYPE(graupel_parameters) :: params + +CONTAINS + + SUBROUTINE graupel_init( & + ccsrim, ccsagg, ccsdep, ccsvel, ccsvxp, ccslam, & + ccslxp, ccsaxp, ccsdxp, ccshi1, ccdvtp, ccidep, & + ccswxp, zconst, zcev, zbev, zcevxp, zbevxp, & + zvzxp, zvz0r, & + v0snow, & + x13o8, x1o2, x27o16, x3o4, x7o4, x7o8, & + zbvi, zcac, zccau, zciau, zcicri, & + zcrcri, zcrfrz, zcrfrz1, zcrfrz2, zeps, zkcac, & + zkphi1, zkphi2, zkphi3, zmi0, zmimax, zmsmin, & + zn0s0, zn0s1, zn0s2, znimax_thom, zqmin, & + zrho0, zthet, zthn, ztmix, ztrfrz, & + zvz0i, icesedi_exp, zams, & + iautocon, isnow_n0temp, dist_cldtop_ref, reduce_dep_ref, & + tmin_iceautoconv, zceff_fac, zceff_min, & + mma_driver, mmb_driver, v_sedi_rain_min, v_sedi_snow_min, v_sedi_graupel_min, & + r_v , & !> gas constant for water vapour + lh_v, & !! latent heat of vapourization + lh_s, & !! latent heat of sublimation + cpdr, & !! (spec. heat of dry air at constant press)^-1 + cvdr, & !! (spec. heat of dry air at const vol)^-1 + b3, & !! melting temperature of ice/snow + t0) !! melting temperature of ice/snow + + INTEGER , INTENT(IN) :: iautocon,isnow_n0temp + REAL(wp), INTENT(IN) :: ccsrim, & + ccsagg, ccsdep, ccsvel, ccsvxp, ccslam, & + ccslxp, ccsaxp, ccsdxp, ccshi1, ccdvtp, ccidep, & + ccswxp, zconst, zcev, zbev, zcevxp, zbevxp, & + zvzxp, zvz0r, & + v0snow, & + x13o8, x1o2, x27o16, x3o4, x7o4, x7o8, & + zbvi, zcac, zccau, zciau, zcicri, & + zcrcri, zcrfrz, zcrfrz1, zcrfrz2, zeps, zkcac, & + zkphi1, zkphi2, zkphi3, zmi0, zmimax, zmsmin, & + zn0s0, zn0s1, zn0s2, znimax_thom, zqmin, & + zrho0, zthet, zthn, ztmix, ztrfrz, & + zvz0i, icesedi_exp, zams, & + dist_cldtop_ref, reduce_dep_ref, & + tmin_iceautoconv, zceff_fac, zceff_min, & + mma_driver(10), mmb_driver(10), v_sedi_rain_min, v_sedi_snow_min, v_sedi_graupel_min, & + r_v , & !> gas constant for water vapour + lh_v, & !! latent heat of vapourization + lh_s, & !! latent heat of sublimation + cpdr, & !! (spec. heat of dry air at constant press)^-1 + cvdr, & !! (spec. heat of dry air at const vol)^-1 + b3, & !! melting temperature of ice/snow + t0 !! melting temperature of ice/snow + + ! gscp_data + params%ccsrim = ccsrim + params%ccsagg = ccsagg + params%ccsdep = ccsdep + params%ccsvel = ccsvel + params%ccsvxp = ccsvxp + params%ccslam = ccslam + params%ccslxp = ccslxp + params%ccsaxp = ccsaxp + params%ccsdxp = ccsdxp + params%ccshi1 = ccshi1 + params%ccdvtp = ccdvtp + params%ccidep = ccidep + params%ccswxp = ccswxp + params%zconst = zconst + params%zcev = zcev + params%zbev = zbev + params%zcevxp = zcevxp + params%zbevxp = zbevxp + params%zvzxp = zvzxp + params%zvz0r = zvz0r + params%v0snow = v0snow + params%x13o8 = x13o8 + params%x1o2 = x1o2 + params%x27o16 = x27o16 + params%x3o4 = x3o4 + params%x7o4 = x7o4 + params%x7o8 = x7o8 + params%zbvi = zbvi + params%zcac = zcac + params%zccau = zccau + params%zciau = zciau + params%zcicri = zcicri + params%zcrcri = zcrcri + params%zcrfrz = zcrfrz + params%zcrfrz1 = zcrfrz1 + params%zcrfrz2 = zcrfrz2 + params%zeps = zeps + params%zkcac = zkcac + params%zkphi1 = zkphi1 + params%zkphi2 = zkphi2 + params%zkphi3 = zkphi3 + params%zmi0 = zmi0 + params%zmimax = zmimax + params%zmsmin = zmsmin + params%zn0s0 = zn0s0 + params%zn0s1 = zn0s1 + params%zn0s2 = zn0s2 + params%znimax_thom = znimax_thom + params%zqmin = zqmin + params%zrho0 = zrho0 + params%zthet = zthet + params%zthn = zthn + params%ztmix = ztmix + params%ztrfrz = ztrfrz + params%zvz0i = zvz0i + params%icesedi_exp = icesedi_exp + params%zams = zams + params%iautocon = iautocon + params%isnow_n0temp = isnow_n0temp + params%dist_cldtop_ref = dist_cldtop_ref + params%reduce_dep_ref = reduce_dep_ref + params%tmin_iceautoconv = tmin_iceautoconv + params%zceff_fac = zceff_fac + params%zceff_min = zceff_min + params%v_sedi_rain_min = v_sedi_rain_min + params%v_sedi_snow_min = v_sedi_snow_min + params%v_sedi_graupel_min = v_sedi_graupel_min + + ! mo_physical constants + params%r_v = r_v + params%lh_v = lh_v + params%lh_s = lh_s + params%cpdr = cpdr + params%cvdr = cvdr + params%b3 = b3 + params%t0 = t0 + + + + mma = mma_driver + mmb = mmb_driver + + !$ACC ENTER DATA COPYIN(mma, mmb) + + END SUBROUTINE graupel_init + + + SUBROUTINE graupel_run( & + nvec,ke, & !> array dimensions + ivstart,ivend, kstart, & !! optional start/end indicies + idbg, & !! optional debug level + zdt, dz, & !! numerics parameters + t,p,rho,qv,qc,qi,qr,qs,qg,qnc, & !! prognostic variables + qi0,qc0, & !! cloud ice/water threshold for autoconversion + & b1, & + & b2w, & + & b4w, & + prr_gsp,prs_gsp,pri_gsp,prg_gsp, & !! surface precipitation rates + qrsflux, & ! total precipitation flux + l_cv, & + ithermo_water, & ! water thermodynamics + ldass_lhn, & + ldiag_ttend, ldiag_qtend , & + ddt_tend_t , ddt_tend_qv , & + ddt_tend_qc , ddt_tend_qi , & !> ddt_tend_xx are tendencies + ddt_tend_qr , ddt_tend_qs)!! necessary for dynamics + + INTEGER, INTENT(IN) :: nvec , & !> number of horizontal points + ke !! number of grid points in vertical direction + + INTEGER, INTENT(IN) :: ivstart , & !> optional start index for horizontal direction + ivend , & !! optional end index for horizontal direction + kstart , & !! optional start index for the vertical index + idbg !! optional debug level + + REAL(KIND=wp), INTENT(IN) :: zdt , & !> time step for integration of microphysics ( s ) + qi0,qc0,& !> cloud ice/water threshold for autoconversion + b1,b2w,b4w + + REAL(KIND=wp), DIMENSION(:,:), INTENT(IN) :: dz , & !> layer thickness of full levels ( m ) + rho , & !! density of moist air (kg/m3) + p !! pressure ( Pa ) + + LOGICAL, INTENT(IN):: l_cv, & !! if true, cv is used instead of cp + ldass_lhn + + INTEGER, INTENT(IN):: ithermo_water !! water thermodynamics + + LOGICAL, INTENT(IN):: ldiag_ttend, & ! if true, temperature tendency shall be diagnosed + ldiag_qtend ! if true, moisture tendencies shall be diagnosed + + REAL(KIND=wp), DIMENSION(:,:), INTENT(INOUT) :: t , & !> temperature ( K ) + qv , & !! specific water vapor content (kg/kg) + qc , & !! specific cloud water content (kg/kg) + qi , & !! specific cloud ice content (kg/kg) + qr , & !! specific rain content (kg/kg) + qs , & !! specific snow content (kg/kg) + qg !! specific graupel content (kg/kg) + + REAL(KIND=wp), INTENT(INOUT) :: qrsflux(:,:) ! total precipitation flux (nudg) + + REAL(KIND=wp), DIMENSION(:), INTENT(INOUT) :: prr_gsp, & !> precipitation rate of rain, grid-scale (kg/(m2*s)) + prs_gsp, & !! precipitation rate of snow, grid-scale (kg/(m2*s)) + prg_gsp, & !! precipitation rate of graupel, grid-scale (kg/(m2*s)) + qnc !! cloud number concentration + + REAL(KIND=wp), DIMENSION(:), INTENT(INOUT):: pri_gsp !! precipitation rate of ice, grid-scale (kg/(m2*s)) + + REAL(KIND=wp), DIMENSION(:,:), INTENT(OUT):: ddt_tend_t , & !> tendency T ( 1/s ) + ddt_tend_qv , & !! tendency qv ( 1/s ) + ddt_tend_qc , & !! tendency qc ( 1/s ) + ddt_tend_qi , & !! tendency qi ( 1/s ) + ddt_tend_qr , & !! tendency qr ( 1/s ) + ddt_tend_qs !! tendency qs ( 1/s ) + + CALL graupel ( & + & nvec =nvec , & !> in: actual array size + & ke =ke , & !< in: actual array size + & ivstart=ivstart , & !< in: start index of calculation + & ivend =ivend , & !< in: end index of calculation + & kstart =kstart , & !< in: vertical start index + & zdt =zdt , & !< in: timestep + & qi0 =qi0 , & + & qc0 =qc0 , & + & b1 = b1, & + & b2w = b2w, & + & b4w = b4w, & + & dz =dz , & !< in: vertical layer thickness + & t =t , & !< in: temp,tracer,... + & p =p , & !< in: full level pres + & rho =rho , & !< in: density + & qv =qv , & !< in: spec. humidity + & qc =qc , & !< in: cloud water + & qi =qi , & !< in: cloud ice + & qr =qr , & !< in: rain water + & qs =qs , & !< in: snow + & qg =qg , & !< in: graupel + & qnc = qnc , & !< cloud number concentration + & prr_gsp=prr_gsp , & !< out: precipitation rate of rain + & prs_gsp=prs_gsp , & !< out: precipitation rate of snow + & pri_gsp=pri_gsp , & !< out: precipitation rate of cloud ice + & prg_gsp=prg_gsp , & !< out: precipitation rate of graupel + & qrsflux= qrsflux , & !< out: precipitation flux + & ldiag_ttend = ldiag_ttend , & !< in: if temp. tendency shall be diagnosed + & ldiag_qtend = ldiag_qtend , & !< in: if moisture tendencies shall be diagnosed + & ddt_tend_t = ddt_tend_t , & !< out: tendency temperature + & ddt_tend_qv = ddt_tend_qv , & !< out: tendency QV + & ddt_tend_qc = ddt_tend_qc , & !< out: tendency QC + & ddt_tend_qi = ddt_tend_qi , & !< out: tendency QI + & ddt_tend_qr = ddt_tend_qr , & !< out: tendency QR + & ddt_tend_qs = ddt_tend_qs , & !< out: tendency QS + & idbg=idbg , & + & l_cv=l_cv , & + & ldass_lhn = ldass_lhn , & + & ithermo_water=ithermo_water, &!< in: latent heat choice + & ccsrim = params%ccsrim, & + & ccsagg = params%ccsagg, & + & ccsdep = params%ccsdep, & + & ccsvel = params%ccsvel, & + & ccsvxp = params%ccsvxp, & + & ccslam = params%ccslam, & + & ccslxp = params%ccslxp, & + & ccsaxp = params%ccsaxp, & + & ccsdxp = params%ccsdxp, & + & ccshi1 = params%ccshi1, & + & ccdvtp = params%ccdvtp, & + & ccidep = params%ccidep, & + & ccswxp = params%ccswxp, & + & zconst = params%zconst, & + & zcev = params%zcev, & + & zbev = params%zbev, & + & zcevxp = params%zcevxp, & + & zbevxp = params%zbevxp, & + & zvzxp = params%zvzxp, & + & zvz0r = params%zvz0r, & + & v0snow = params%v0snow, & + & x13o8 = params%x13o8, & + & x1o2 = params%x1o2, & + & x27o16 = params%x27o16, & + & x3o4 = params%x3o4, & + & x7o4 = params%x7o4, & + & x7o8 = params%x7o8, & + & zbvi = params%zbvi, & + & zcac = params%zcac, & + & zccau = params%zccau, & + & zciau = params%zciau, & + & zcicri = params%zcicri, & + & zcrcri = params%zcrcri, & + & zcrfrz = params%zcrfrz, & + & zcrfrz1 = params%zcrfrz1, & + & zcrfrz2 = params%zcrfrz2, & + & zeps = params%zeps, & + & zkcac = params%zkcac, & + & zkphi1 = params%zkphi1, & + & zkphi2 = params%zkphi2, & + & zkphi3 = params%zkphi3, & + & zmi0 = params%zmi0, & + & zmimax = params%zmimax, & + & zmsmin = params%zmsmin, & + & zn0s0 = params%zn0s0, & + & zn0s1 = params%zn0s1, & + & zn0s2 = params%zn0s2, & + & znimax_thom = params%znimax_thom, & + & zqmin = params%zqmin, & + & zrho0 = params%zrho0, & + & zthet = params%zthet, & + & zthn = params%zthn, & + & ztmix = params%ztmix, & + & ztrfrz = params%ztrfrz, & + & zvz0i = params%zvz0i, & + & icesedi_exp = params%icesedi_exp, & + & zams = params%zams, & + & iautocon = params%iautocon, & + & isnow_n0temp = params%isnow_n0temp, & + & dist_cldtop_ref = params%dist_cldtop_ref, & + & reduce_dep_ref = params%reduce_dep_ref, & + & tmin_iceautoconv = params%tmin_iceautoconv, & + & zceff_fac = params%zceff_fac, & + & zceff_min = params%zceff_min, & + & mma = mma, & + & mmb = mmb, & + & v_sedi_rain_min = params%v_sedi_rain_min, & + & v_sedi_snow_min = params%v_sedi_snow_min, & + & v_sedi_graupel_min = params%v_sedi_graupel_min, & + & r_v = params%r_v, & + & lh_v = params%lh_v, & + & lh_s = params%lh_s, & + & cpdr = params%cpdr, & + & cvdr = params%cvdr, & + & b3 = params%b3, & + t0 = params%t0) + + END SUBROUTINE graupel_run + +END MODULE mo_graupel_granule + diff --git a/tools/tests/f2ser/test_parsing.py b/tools/tests/f2ser/test_parsing.py index cf06d81e00..3d5858def8 100644 --- a/tools/tests/f2ser/test_parsing.py +++ b/tools/tests/f2ser/test_parsing.py @@ -58,3 +58,16 @@ def test_granule_parsing_no_intent(samples_path): parser = GranuleParser(samples_path / "subroutine_example.f90", []) with pytest.raises(ParsingError): parser() + + +def test_multiline_declaration_parsing(samples_path): + parser = GranuleParser(samples_path / "multiline_example.f90", []) + parsed_granule = parser() + subroutines = parsed_granule.subroutines + assert list(subroutines) == ["graupel_init", "graupel_run"] + assert subroutines["graupel_init"]["in"]["codegen_ctx"] == CodegenContext( + first_declaration_ln=121, last_declaration_ln=145, end_subroutine_ln=231 + ) + assert subroutines["graupel_run"]["in"]["codegen_ctx"] == CodegenContext( + first_declaration_ln=254, last_declaration_ln=301, end_subroutine_ln=419 + ) diff --git a/tools/tests/icon4pygen/helpers.py b/tools/tests/icon4pygen/helpers.py index 064e98e060..82ca300551 100644 --- a/tools/tests/icon4pygen/helpers.py +++ b/tools/tests/icon4pygen/helpers.py @@ -11,6 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from icon4pytools.common import ICON4PY_MODEL_QUALIFIED_NAME + def get_stencil_module_path(stencil_module: str, stencil_name: str) -> str: - return f"icon4py.{stencil_module}.{stencil_name}:{stencil_name}" + return f"{ICON4PY_MODEL_QUALIFIED_NAME}.{stencil_module}.{stencil_name}:{stencil_name}" diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index de476bb761..898b6b71a0 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -15,6 +15,7 @@ import pkgutil import re +import icon4py.model.atmosphere.dycore as dycore import pytest from click.testing import CliRunner from icon4pytools.icon4pygen.cli import main @@ -24,6 +25,7 @@ from .helpers import get_stencil_module_path +DYCORE_PKG = "atmosphere.dycore" LEVELS_PER_THREAD = "1" BLOCK_SIZE = "128" OUTPATH = "." @@ -34,10 +36,10 @@ def cli(): return CliRunner() -def atm_dyn_iconam_fencils() -> list[tuple[str, str]]: - pkgpath = os.path.dirname(icon4py.atm_dyn_iconam.__file__) +def dycore_fencils() -> list[tuple[str, str]]: + pkgpath = os.path.dirname(dycore.__file__) stencils = [name for _, name, _ in pkgutil.iter_modules([pkgpath])] - fencils = [("atm_dyn_iconam", stencil) for stencil in stencils] + fencils = [(DYCORE_PKG, stencil) for stencil in stencils] return fencils @@ -110,8 +112,8 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") -@pytest.mark.parametrize(("stencil_module", "stencil_name"), atm_dyn_iconam_fencils()) -def test_codegen_atm_dyn_iconam(cli, stencil_module, stencil_name) -> None: +@pytest.mark.parametrize(("stencil_module", "stencil_name"), dycore_fencils()) +def test_codegen_dycore(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) with cli.isolated_filesystem(): result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) diff --git a/tools/tests/icon4pygen/test_exceptions.py b/tools/tests/icon4pygen/test_exceptions.py index 69445d5881..d86a3854e8 100644 --- a/tools/tests/icon4pygen/test_exceptions.py +++ b/tools/tests/icon4pygen/test_exceptions.py @@ -14,7 +14,7 @@ import pytest from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Dimension, Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim from icon4pytools.icon4pygen.bindings.codegen.render.location import LocationRenderer from icon4pytools.icon4pygen.bindings.entities import Offset, chain_from_str diff --git a/tools/tests/icon4pygen/test_field_rendering.py b/tools/tests/icon4pygen/test_field_rendering.py index f6649b7c62..400b8aa651 100644 --- a/tools/tests/icon4pygen/test_field_rendering.py +++ b/tools/tests/icon4pygen/test_field_rendering.py @@ -13,10 +13,11 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim + from icon4pytools.icon4pygen.bindings.workflow import PyBindGen from icon4pytools.icon4pygen.metadata import get_stencil_info -from icon4py.common.dimension import E2CDim, EdgeDim, KDim def test_horizontal_field_sid_rendering(): diff --git a/tools/tests/icon4pygen/test_metadata.py b/tools/tests/icon4pygen/test_metadata.py index c0b50f9751..c23665a963 100644 --- a/tools/tests/icon4pygen/test_metadata.py +++ b/tools/tests/icon4pygen/test_metadata.py @@ -14,7 +14,7 @@ import pytest from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim from icon4pytools.icon4pygen.metadata import _get_field_infos, provide_neighbor_table diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 1e0de1b5ca..e109dcaefa 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -40,7 +40,7 @@ def test_stencil_collector_invalid_module(): def test_stencil_collector_invalid_member(): - from icon4py.atm_dyn_iconam import apply_nabla2_to_w + from icon4py.model.atmosphere.dycore import apply_nabla2_to_w module_path = Path(apply_nabla2_to_w.__file__) parents = module_path.parents[0] diff --git a/tools/tox.ini b/tools/tox.ini index f645390a9b..0907e9a140 100644 --- a/tools/tox.ini +++ b/tools/tox.ini @@ -26,6 +26,8 @@ allowlist_externals = [testenv:dev] basepython = python3.10 +setenv = + PIP_SRC = _external_src skip_install = true commands = commands_post = diff --git a/tox.ini b/tox.ini index 1f63dc4c2c..65d9765d16 100644 --- a/tox.ini +++ b/tox.ini @@ -11,22 +11,23 @@ skipsdist = true passenv = PIP_USER PYTHONUSERBASE -setenv = - PIP_SRC = _external_src deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules advection/src atm_dyn_iconam/src common/src pyutils/src testutils/src - pytest -v -s -n auto --cov --cov-append --ignore=tools/ + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/common/src tools/src + pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = - rm -Rf _reports/coverage_html - coverage html + rm -rf tests/_reports/coverage_html + -coverage html + -coverage report allowlist_externals = /bin/bash rm [testenv:dev] basepython = python3.10 +setenv = + PIP_SRC = _external_src skip_install = true commands = commands_post = From 2322c5ee8c9f86aafcfc41729e7d2a253a95644d Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 27 Jul 2023 12:34:24 +0200 Subject: [PATCH 083/105] Moved files according to icon4py model structure --- .../advection/src/icon4py/model/atmosphere}/advection/__init__.py | 0 .../icon4py/model/atmosphere}/advection/btraj_dreg_stencil_01.py | 0 .../icon4py/model/atmosphere}/advection/btraj_dreg_stencil_02.py | 0 .../icon4py/model/atmosphere}/advection/btraj_dreg_stencil_03.py | 0 .../atmosphere}/advection/divide_flux_area_list_stencil_01.py | 0 .../atmosphere}/advection/divide_flux_area_list_stencil_02.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_01.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_02.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_02a.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_02b.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_02c.py | 0 .../model/atmosphere}/advection/face_val_ppm_stencil_05.py | 0 .../model/atmosphere}/advection/hflux_ffsl_hybrid_stencil_01a.py | 0 .../model/atmosphere}/advection/hflux_ffsl_hybrid_stencil_02.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_01a.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_01b.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_02.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_03.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_04.py | 0 .../model/atmosphere}/advection/hflx_limiter_mo_stencil_05.py | 0 .../model/atmosphere}/advection/hflx_limiter_pd_stencil_01.py | 0 .../model/atmosphere}/advection/hflx_limiter_pd_stencil_02.py | 0 .../src/icon4py/model/atmosphere}/advection/hor_adv_stencil_01.py | 0 .../atmosphere}/advection/prep_gauss_quadrature_c_list_stencil.py | 0 .../atmosphere}/advection/prep_gauss_quadrature_c_stencil.py | 0 .../model/atmosphere}/advection/rbf_intp_edge_stencil_01.py | 0 .../model/atmosphere}/advection/recon_lsq_cell_c_stencil.py | 0 .../model/atmosphere}/advection/recon_lsq_cell_c_svd_stencil.py | 0 .../model/atmosphere}/advection/recon_lsq_cell_l_svd_stencil.py | 0 .../src/icon4py/model/atmosphere}/advection/set_zero_c.py | 0 .../src/icon4py/model/atmosphere}/advection/set_zero_c_k.py | 0 .../model/atmosphere}/advection/step_advection_stencil_01.py | 0 .../model/atmosphere}/advection/step_advection_stencil_02.py | 0 .../model/atmosphere}/advection/step_advection_stencil_03.py | 0 .../model/atmosphere}/advection/step_advection_stencil_04.py | 0 .../model/atmosphere}/advection/upwind_hflux_miura3_stencil_01.py | 0 .../atmosphere}/advection/upwind_hflux_miura_cycl_stencil_01.py | 0 .../atmosphere}/advection/upwind_hflux_miura_cycl_stencil_02.py | 0 .../atmosphere}/advection/upwind_hflux_miura_cycl_stencil_03a.py | 0 .../atmosphere}/advection/upwind_hflux_miura_cycl_stencil_03b.py | 0 .../model/atmosphere}/advection/upwind_hflux_miura_stencil_01.py | 0 .../model/atmosphere}/advection/upwind_vflux_ppm_stencil_01.py | 0 .../model/atmosphere}/advection/v_limit_prbl_sm_stencil_01.py | 0 .../model/atmosphere}/advection/v_limit_prbl_sm_stencil_02.py | 0 .../icon4py/model/atmosphere}/advection/vert_adv_stencil_01.py | 0 {advection => model/atmosphere/advection}/tests/__init__.py | 0 .../atmosphere/advection}/tests/test_btraj_dreg_stencil_01.py | 0 .../atmosphere/advection}/tests/test_btraj_dreg_stencil_02.py | 0 .../atmosphere/advection}/tests/test_btraj_dreg_stencil_03.py | 0 .../advection}/tests/test_divide_flux_area_list_stencil_01.py | 0 .../advection}/tests/test_divide_flux_area_list_stencil_02.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_01.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_02.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_02a.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_02b.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_02c.py | 0 .../atmosphere/advection}/tests/test_face_val_ppm_stencil_05.py | 0 .../advection}/tests/test_hflux_ffsl_hybrid_stencil_01a.py | 0 .../advection}/tests/test_hflux_ffsl_hybrid_stencil_02.py | 0 .../advection}/tests/test_hflx_limiter_mo_stencil_01a.py | 0 .../advection}/tests/test_hflx_limiter_mo_stencil_01b.py | 0 .../advection}/tests/test_hflx_limiter_mo_stencil_02.py | 0 .../advection}/tests/test_hflx_limiter_mo_stencil_03.py | 0 .../advection}/tests/test_hflx_limiter_mo_stencil_04.py | 0 .../advection}/tests/test_hflx_limiter_pd_stencil_01.py | 0 .../advection}/tests/test_hflx_limiter_pd_stencil_02.py | 0 .../atmosphere/advection}/tests/test_hor_adv_stencil_01.py | 0 .../advection}/tests/test_prep_gauss_quadrature_c_list_stencil.py | 0 .../advection}/tests/test_prep_gauss_quadrature_c_stencil.py | 0 .../atmosphere/advection}/tests/test_rbf_intp_edge_stencil_01.py | 0 .../atmosphere/advection}/tests/test_recon_lsq_cell_c_stencil.py | 0 .../advection}/tests/test_recon_lsq_cell_c_svd_stencil.py | 0 .../advection}/tests/test_recon_lsq_cell_l_svd_stencil.py | 0 .../atmosphere/advection}/tests/test_set_zero_c.py | 0 .../atmosphere/advection}/tests/test_set_zero_c_k.py | 0 .../atmosphere/advection}/tests/test_step_advection_stencil_01.py | 0 .../atmosphere/advection}/tests/test_step_advection_stencil_02.py | 0 .../atmosphere/advection}/tests/test_step_advection_stencil_03.py | 0 .../atmosphere/advection}/tests/test_step_advection_stencil_04.py | 0 .../advection}/tests/test_upwind_hflux_miura3_stencil_01.py | 0 .../advection}/tests/test_upwind_hflux_miura_cycl_stencil_01.py | 0 .../advection}/tests/test_upwind_hflux_miura_cycl_stencil_02.py | 0 .../advection}/tests/test_upwind_hflux_miura_cycl_stencil_03a.py | 0 .../advection}/tests/test_upwind_hflux_miura_cycl_stencil_03b.py | 0 .../advection}/tests/test_upwind_hflux_miura_stencil_01.py | 0 .../advection}/tests/test_upwind_vflux_ppm_stencil_01.py | 0 .../atmosphere/advection}/tests/test_utils/__init__.py | 0 .../atmosphere/advection}/tests/test_utils/helpers.py | 0 .../atmosphere/advection}/tests/test_utils/simple_mesh.py | 0 .../atmosphere/advection}/tests/test_vert_adv_stencil_01.py | 0 .../atmosphere/advection}/tests/test_vlimit_prbl_sm_stencil_01.py | 0 .../atmosphere/advection}/tests/test_vlimit_prbl_sm_stencil_02.py | 0 92 files changed, 0 insertions(+), 0 deletions(-) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/__init__.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/btraj_dreg_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/btraj_dreg_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/btraj_dreg_stencil_03.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/divide_flux_area_list_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/divide_flux_area_list_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_02a.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_02b.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_02c.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/face_val_ppm_stencil_05.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflux_ffsl_hybrid_stencil_01a.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflux_ffsl_hybrid_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_01a.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_01b.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_03.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_04.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_mo_stencil_05.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_pd_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hflx_limiter_pd_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/hor_adv_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/prep_gauss_quadrature_c_list_stencil.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/prep_gauss_quadrature_c_stencil.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/rbf_intp_edge_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/recon_lsq_cell_c_stencil.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/recon_lsq_cell_c_svd_stencil.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/recon_lsq_cell_l_svd_stencil.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/set_zero_c.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/set_zero_c_k.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/step_advection_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/step_advection_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/step_advection_stencil_03.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/step_advection_stencil_04.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura3_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura_cycl_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura_cycl_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura_cycl_stencil_03a.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura_cycl_stencil_03b.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_hflux_miura_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/upwind_vflux_ppm_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/v_limit_prbl_sm_stencil_01.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/v_limit_prbl_sm_stencil_02.py (100%) rename {advection/src/icon4py => model/atmosphere/advection/src/icon4py/model/atmosphere}/advection/vert_adv_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/__init__.py (100%) rename {advection => model/atmosphere/advection}/tests/test_btraj_dreg_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_btraj_dreg_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_btraj_dreg_stencil_03.py (100%) rename {advection => model/atmosphere/advection}/tests/test_divide_flux_area_list_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_divide_flux_area_list_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_02a.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_02b.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_02c.py (100%) rename {advection => model/atmosphere/advection}/tests/test_face_val_ppm_stencil_05.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflux_ffsl_hybrid_stencil_01a.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflux_ffsl_hybrid_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_mo_stencil_01a.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_mo_stencil_01b.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_mo_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_mo_stencil_03.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_mo_stencil_04.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_pd_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hflx_limiter_pd_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_hor_adv_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_prep_gauss_quadrature_c_list_stencil.py (100%) rename {advection => model/atmosphere/advection}/tests/test_prep_gauss_quadrature_c_stencil.py (100%) rename {advection => model/atmosphere/advection}/tests/test_rbf_intp_edge_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_recon_lsq_cell_c_stencil.py (100%) rename {advection => model/atmosphere/advection}/tests/test_recon_lsq_cell_c_svd_stencil.py (100%) rename {advection => model/atmosphere/advection}/tests/test_recon_lsq_cell_l_svd_stencil.py (100%) rename {advection => model/atmosphere/advection}/tests/test_set_zero_c.py (100%) rename {advection => model/atmosphere/advection}/tests/test_set_zero_c_k.py (100%) rename {advection => model/atmosphere/advection}/tests/test_step_advection_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_step_advection_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_step_advection_stencil_03.py (100%) rename {advection => model/atmosphere/advection}/tests/test_step_advection_stencil_04.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura3_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura_cycl_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura_cycl_stencil_02.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura_cycl_stencil_03a.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura_cycl_stencil_03b.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_hflux_miura_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_upwind_vflux_ppm_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_utils/__init__.py (100%) rename {advection => model/atmosphere/advection}/tests/test_utils/helpers.py (100%) rename {advection => model/atmosphere/advection}/tests/test_utils/simple_mesh.py (100%) rename {advection => model/atmosphere/advection}/tests/test_vert_adv_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_vlimit_prbl_sm_stencil_01.py (100%) rename {advection => model/atmosphere/advection}/tests/test_vlimit_prbl_sm_stencil_02.py (100%) diff --git a/advection/src/icon4py/advection/__init__.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py similarity index 100% rename from advection/src/icon4py/advection/__init__.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/btraj_dreg_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/btraj_dreg_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py diff --git a/advection/src/icon4py/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py similarity index 100% rename from advection/src/icon4py/advection/btraj_dreg_stencil_03.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/divide_flux_area_list_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py diff --git a/advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/divide_flux_area_list_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_02a.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_02b.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_02c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_02c.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py diff --git a/advection/src/icon4py/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py similarity index 100% rename from advection/src/icon4py/advection/face_val_ppm_stencil_05.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py similarity index 100% rename from advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_01a.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py diff --git a/advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/hflux_ffsl_hybrid_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_01a.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_01b.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_mo_stencil_05.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py diff --git a/advection/src/icon4py/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/hor_adv_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py similarity index 100% rename from advection/src/icon4py/advection/prep_gauss_quadrature_c_list_stencil.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py diff --git a/advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py similarity index 100% rename from advection/src/icon4py/advection/prep_gauss_quadrature_c_stencil.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py diff --git a/advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py similarity index 100% rename from advection/src/icon4py/advection/recon_lsq_cell_c_stencil.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py diff --git a/advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py similarity index 100% rename from advection/src/icon4py/advection/recon_lsq_cell_c_svd_stencil.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py diff --git a/advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py similarity index 100% rename from advection/src/icon4py/advection/recon_lsq_cell_l_svd_stencil.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py diff --git a/advection/src/icon4py/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py similarity index 100% rename from advection/src/icon4py/advection/set_zero_c.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py diff --git a/advection/src/icon4py/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py similarity index 100% rename from advection/src/icon4py/advection/set_zero_c_k.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py diff --git a/advection/src/icon4py/advection/step_advection_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/step_advection_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py diff --git a/advection/src/icon4py/advection/step_advection_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/step_advection_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py diff --git a/advection/src/icon4py/advection/step_advection_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py similarity index 100% rename from advection/src/icon4py/advection/step_advection_stencil_03.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py diff --git a/advection/src/icon4py/advection/step_advection_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py similarity index 100% rename from advection/src/icon4py/advection/step_advection_stencil_04.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura3_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03a.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura_cycl_stencil_03b.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py diff --git a/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py diff --git a/advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py similarity index 100% rename from advection/src/icon4py/advection/v_limit_prbl_sm_stencil_02.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py diff --git a/advection/src/icon4py/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py similarity index 100% rename from advection/src/icon4py/advection/vert_adv_stencil_01.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py diff --git a/advection/tests/__init__.py b/model/atmosphere/advection/tests/__init__.py similarity index 100% rename from advection/tests/__init__.py rename to model/atmosphere/advection/tests/__init__.py diff --git a/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py similarity index 100% rename from advection/tests/test_btraj_dreg_stencil_01.py rename to model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py diff --git a/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py similarity index 100% rename from advection/tests/test_btraj_dreg_stencil_02.py rename to model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py diff --git a/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py similarity index 100% rename from advection/tests/test_btraj_dreg_stencil_03.py rename to model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py diff --git a/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py similarity index 100% rename from advection/tests/test_divide_flux_area_list_stencil_01.py rename to model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py diff --git a/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py similarity index 100% rename from advection/tests/test_divide_flux_area_list_stencil_02.py rename to model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py diff --git a/advection/tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_01.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py diff --git a/advection/tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_02.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py diff --git a/advection/tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_02a.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py diff --git a/advection/tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_02b.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py diff --git a/advection/tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_02c.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py diff --git a/advection/tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py similarity index 100% rename from advection/tests/test_face_val_ppm_stencil_05.py rename to model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py diff --git a/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py similarity index 100% rename from advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py rename to model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py diff --git a/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py similarity index 100% rename from advection/tests/test_hflux_ffsl_hybrid_stencil_02.py rename to model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py diff --git a/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py similarity index 100% rename from advection/tests/test_hflx_limiter_mo_stencil_01a.py rename to model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py diff --git a/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py similarity index 100% rename from advection/tests/test_hflx_limiter_mo_stencil_01b.py rename to model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py similarity index 100% rename from advection/tests/test_hflx_limiter_mo_stencil_02.py rename to model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py similarity index 100% rename from advection/tests/test_hflx_limiter_mo_stencil_03.py rename to model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py similarity index 100% rename from advection/tests/test_hflx_limiter_mo_stencil_04.py rename to model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py diff --git a/advection/tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py similarity index 100% rename from advection/tests/test_hflx_limiter_pd_stencil_01.py rename to model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py similarity index 100% rename from advection/tests/test_hflx_limiter_pd_stencil_02.py rename to model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py diff --git a/advection/tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py similarity index 100% rename from advection/tests/test_hor_adv_stencil_01.py rename to model/atmosphere/advection/tests/test_hor_adv_stencil_01.py diff --git a/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py similarity index 100% rename from advection/tests/test_prep_gauss_quadrature_c_list_stencil.py rename to model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py diff --git a/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py similarity index 100% rename from advection/tests/test_prep_gauss_quadrature_c_stencil.py rename to model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py diff --git a/advection/tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py similarity index 100% rename from advection/tests/test_rbf_intp_edge_stencil_01.py rename to model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py diff --git a/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py similarity index 100% rename from advection/tests/test_recon_lsq_cell_c_stencil.py rename to model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py diff --git a/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py similarity index 100% rename from advection/tests/test_recon_lsq_cell_c_svd_stencil.py rename to model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py diff --git a/advection/tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py similarity index 100% rename from advection/tests/test_recon_lsq_cell_l_svd_stencil.py rename to model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py diff --git a/advection/tests/test_set_zero_c.py b/model/atmosphere/advection/tests/test_set_zero_c.py similarity index 100% rename from advection/tests/test_set_zero_c.py rename to model/atmosphere/advection/tests/test_set_zero_c.py diff --git a/advection/tests/test_set_zero_c_k.py b/model/atmosphere/advection/tests/test_set_zero_c_k.py similarity index 100% rename from advection/tests/test_set_zero_c_k.py rename to model/atmosphere/advection/tests/test_set_zero_c_k.py diff --git a/advection/tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py similarity index 100% rename from advection/tests/test_step_advection_stencil_01.py rename to model/atmosphere/advection/tests/test_step_advection_stencil_01.py diff --git a/advection/tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py similarity index 100% rename from advection/tests/test_step_advection_stencil_02.py rename to model/atmosphere/advection/tests/test_step_advection_stencil_02.py diff --git a/advection/tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py similarity index 100% rename from advection/tests/test_step_advection_stencil_03.py rename to model/atmosphere/advection/tests/test_step_advection_stencil_03.py diff --git a/advection/tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py similarity index 100% rename from advection/tests/test_step_advection_stencil_04.py rename to model/atmosphere/advection/tests/test_step_advection_stencil_04.py diff --git a/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura3_stencil_01.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py diff --git a/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py diff --git a/advection/tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py similarity index 100% rename from advection/tests/test_upwind_hflux_miura_stencil_01.py rename to model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py similarity index 100% rename from advection/tests/test_upwind_vflux_ppm_stencil_01.py rename to model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py diff --git a/advection/tests/test_utils/__init__.py b/model/atmosphere/advection/tests/test_utils/__init__.py similarity index 100% rename from advection/tests/test_utils/__init__.py rename to model/atmosphere/advection/tests/test_utils/__init__.py diff --git a/advection/tests/test_utils/helpers.py b/model/atmosphere/advection/tests/test_utils/helpers.py similarity index 100% rename from advection/tests/test_utils/helpers.py rename to model/atmosphere/advection/tests/test_utils/helpers.py diff --git a/advection/tests/test_utils/simple_mesh.py b/model/atmosphere/advection/tests/test_utils/simple_mesh.py similarity index 100% rename from advection/tests/test_utils/simple_mesh.py rename to model/atmosphere/advection/tests/test_utils/simple_mesh.py diff --git a/advection/tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py similarity index 100% rename from advection/tests/test_vert_adv_stencil_01.py rename to model/atmosphere/advection/tests/test_vert_adv_stencil_01.py diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py similarity index 100% rename from advection/tests/test_vlimit_prbl_sm_stencil_01.py rename to model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py diff --git a/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py similarity index 100% rename from advection/tests/test_vlimit_prbl_sm_stencil_02.py rename to model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py From 55688dd10f77e2b290114a7ccabd5d0f4d0ce666 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 27 Jul 2023 12:36:24 +0200 Subject: [PATCH 084/105] Add README and requirements to advection folder --- model/atmosphere/advection/README.md | 9 ++ model/atmosphere/advection/pyproject.toml | 118 ++++++++++++++++++ .../atmosphere/advection/requirements-dev.txt | 3 + model/atmosphere/advection/requirements.txt | 3 + 4 files changed, 133 insertions(+) create mode 100644 model/atmosphere/advection/README.md create mode 100644 model/atmosphere/advection/pyproject.toml create mode 100644 model/atmosphere/advection/requirements-dev.txt create mode 100644 model/atmosphere/advection/requirements.txt diff --git a/model/atmosphere/advection/README.md b/model/atmosphere/advection/README.md new file mode 100644 index 0000000000..524a73d07a --- /dev/null +++ b/model/atmosphere/advection/README.md @@ -0,0 +1,9 @@ +# icon4py-atmosphere-advection + +## Description + +Contains code ported from ICON `src/advection`, which is the advection component of the ICON model. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml new file mode 100644 index 0000000000..268c47eed0 --- /dev/null +++ b/model/atmosphere/advection/pyproject.toml @@ -0,0 +1,118 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel>=0.40.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "icon4py-atmosphere-advection" +description = "ICON advection" +readme = "README.md" +requires-python = ">=3.10" +license = {file = "LICENSE"} +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "icon4py_common>=0.0.5", +] +dynamic = ['version'] + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +parallel = true +branch = true +source_pkgs = ['advection'] + +[tool.isort] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +multi_line_output = 3 +use_parentheses = true +include_trailing_comma = true +force_grid_wrap = 0 + +[tool.mypy] +install_types = true +non_interactive = true +exclude = [ + '^tests/*.py', +] +disallow_incomplete_defs = true +disallow_untyped_defs = true +ignore_missing_imports = false +implicit_reexport = true +warn_unused_configs = true +warn_unused_ignores = true +warn_redundant_casts = true +show_column_numbers = true +show_error_codes = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.atmosphere.advection' = ['py.typed'] diff --git a/model/atmosphere/advection/requirements-dev.txt b/model/atmosphere/advection/requirements-dev.txt new file mode 100644 index 0000000000..32d385e08d --- /dev/null +++ b/model/atmosphere/advection/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements-dev.txt +-e ../../common +-e . diff --git a/model/atmosphere/advection/requirements.txt b/model/atmosphere/advection/requirements.txt new file mode 100644 index 0000000000..79787ec137 --- /dev/null +++ b/model/atmosphere/advection/requirements.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements.txt +../../common +. From cded6f55392ccdd226972d8e68e47c1c4b9a5422 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 28 Jul 2023 13:21:15 +0200 Subject: [PATCH 085/105] Adapt to icon4py refactor for advection --- advection/README.md | 9 --- advection/__init__.py | 12 ---- advection/requirements-dev.txt | 3 - advection/requirements.txt | 3 - advection/setup.cfg | 50 ----------------- model/README.md | 2 + .../model/atmosphere/advection/__init__.py | 21 +++++++ .../advection/btraj_dreg_stencil_01.py | 2 +- .../advection/btraj_dreg_stencil_02.py | 4 +- .../advection/btraj_dreg_stencil_03.py | 2 +- .../divide_flux_area_list_stencil_01.py | 4 +- .../divide_flux_area_list_stencil_02.py | 2 +- .../advection/face_val_ppm_stencil_01.py | 2 +- .../advection/face_val_ppm_stencil_02.py | 2 +- .../advection/face_val_ppm_stencil_02a.py | 2 +- .../advection/face_val_ppm_stencil_02b.py | 2 +- .../advection/face_val_ppm_stencil_02c.py | 2 +- .../advection/face_val_ppm_stencil_05.py | 2 +- .../hflux_ffsl_hybrid_stencil_01a.py | 2 +- .../advection/hflux_ffsl_hybrid_stencil_02.py | 2 +- .../advection/hflx_limiter_mo_stencil_01a.py | 2 +- .../advection/hflx_limiter_mo_stencil_01b.py | 2 +- .../advection/hflx_limiter_mo_stencil_02.py | 2 +- .../advection/hflx_limiter_mo_stencil_03.py | 2 +- .../advection/hflx_limiter_mo_stencil_04.py | 2 +- .../advection/hflx_limiter_mo_stencil_05.py | 2 +- .../advection/hflx_limiter_pd_stencil_01.py | 2 +- .../advection/hflx_limiter_pd_stencil_02.py | 2 +- .../advection/hor_adv_stencil_01.py | 2 +- .../prep_gauss_quadrature_c_list_stencil.py | 42 +++++++------- .../prep_gauss_quadrature_c_stencil.py | 50 ++++++++--------- .../advection/rbf_intp_edge_stencil_01.py | 2 +- .../advection/recon_lsq_cell_c_stencil.py | 2 +- .../advection/recon_lsq_cell_c_svd_stencil.py | 2 +- .../advection/recon_lsq_cell_l_svd_stencil.py | 2 +- .../model/atmosphere/advection/set_zero_c.py | 2 +- .../atmosphere/advection/set_zero_c_k.py | 2 +- .../advection/step_advection_stencil_01.py | 2 +- .../advection/step_advection_stencil_02.py | 2 +- .../advection/step_advection_stencil_03.py | 2 +- .../advection/step_advection_stencil_04.py | 2 +- .../upwind_hflux_miura3_stencil_01.py | 2 +- .../upwind_hflux_miura_cycl_stencil_01.py | 2 +- .../upwind_hflux_miura_cycl_stencil_02.py | 2 +- .../upwind_hflux_miura_cycl_stencil_03a.py | 2 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 2 +- .../upwind_hflux_miura_stencil_01.py | 2 +- .../advection/upwind_vflux_ppm_stencil_01.py | 2 +- .../advection/v_limit_prbl_sm_stencil_01.py | 2 +- .../advection/v_limit_prbl_sm_stencil_02.py | 2 +- .../advection/vert_adv_stencil_01.py | 2 +- .../atmosphere/advection/tests/conftest.py | 21 ++++++- .../tests/test_btraj_dreg_stencil_01.py | 8 +-- .../tests/test_btraj_dreg_stencil_02.py | 18 +++--- .../tests/test_btraj_dreg_stencil_03.py | 8 +-- .../test_divide_flux_area_list_stencil_01.py | 10 ++-- .../test_divide_flux_area_list_stencil_02.py | 8 +-- .../tests/test_face_val_ppm_stencil_01.py | 8 +-- .../tests/test_face_val_ppm_stencil_02.py | 8 +-- .../tests/test_face_val_ppm_stencil_02a.py | 8 +-- .../tests/test_face_val_ppm_stencil_02b.py | 8 +-- .../tests/test_face_val_ppm_stencil_02c.py | 8 +-- .../tests/test_face_val_ppm_stencil_05.py | 8 +-- .../test_hflux_ffsl_hybrid_stencil_01a.py | 8 +-- .../test_hflux_ffsl_hybrid_stencil_02.py | 8 +-- .../tests/test_hflx_limiter_mo_stencil_01a.py | 8 +-- .../tests/test_hflx_limiter_mo_stencil_01b.py | 8 +-- .../tests/test_hflx_limiter_mo_stencil_02.py | 8 +-- .../tests/test_hflx_limiter_mo_stencil_03.py | 8 +-- .../tests/test_hflx_limiter_mo_stencil_04.py | 8 +-- .../tests/test_hflx_limiter_pd_stencil_01.py | 8 +-- .../tests/test_hflx_limiter_pd_stencil_02.py | 8 +-- .../tests/test_hor_adv_stencil_01.py | 8 +-- ...st_prep_gauss_quadrature_c_list_stencil.py | 48 ++++++++-------- .../test_prep_gauss_quadrature_c_stencil.py | 56 +++++++++---------- .../tests/test_rbf_intp_edge_stencil_01.py | 8 +-- .../tests/test_recon_lsq_cell_c_stencil.py | 8 +-- .../test_recon_lsq_cell_c_svd_stencil.py | 8 +-- .../test_recon_lsq_cell_l_svd_stencil.py | 8 +-- .../advection/tests/test_set_zero_c.py | 8 +-- .../advection/tests/test_set_zero_c_k.py | 8 +-- .../tests/test_step_advection_stencil_01.py | 8 +-- .../tests/test_step_advection_stencil_02.py | 8 +-- .../tests/test_step_advection_stencil_03.py | 8 +-- .../tests/test_step_advection_stencil_04.py | 8 +-- .../test_upwind_hflux_miura3_stencil_01.py | 8 +-- ...test_upwind_hflux_miura_cycl_stencil_01.py | 8 +-- ...test_upwind_hflux_miura_cycl_stencil_02.py | 8 +-- ...est_upwind_hflux_miura_cycl_stencil_03a.py | 8 +-- ...est_upwind_hflux_miura_cycl_stencil_03b.py | 8 +-- .../test_upwind_hflux_miura_stencil_01.py | 8 +-- .../tests/test_upwind_vflux_ppm_stencil_01.py | 8 +-- .../tests/test_vert_adv_stencil_01.py | 8 +-- .../tests/test_vlimit_prbl_sm_stencil_01.py | 8 +-- .../tests/test_vlimit_prbl_sm_stencil_02.py | 8 +-- .../atmosphere/dycore}/compute_airmass.py | 2 +- .../dycore}/tests/test_compute_airmass.py | 0 .../model/common/test_utils/simple_mesh.py | 31 ++++++++++ model/requirements-dev.txt | 1 + model/requirements.txt | 1 + model/tox.ini | 1 + .../src/icon4pytools/liskov/external/gt4py.py | 2 +- tools/tests/icon4pygen/test_codegen.py | 14 +++++ tox.ini | 1 + 104 files changed, 404 insertions(+), 394 deletions(-) delete mode 100644 advection/README.md delete mode 100644 advection/__init__.py delete mode 100644 advection/requirements-dev.txt delete mode 100644 advection/requirements.txt delete mode 100644 advection/setup.cfg rename advection/setup.py => model/atmosphere/advection/tests/conftest.py (54%) rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore}/compute_airmass.py (95%) rename {atm_dyn_iconam => model/atmosphere/dycore}/tests/test_compute_airmass.py (100%) diff --git a/advection/README.md b/advection/README.md deleted file mode 100644 index 9ee5f3fdab..0000000000 --- a/advection/README.md +++ /dev/null @@ -1,9 +0,0 @@ -# icon4py-advection - -## Description - -Contains code ported from ICON `src/advection`. - -## Installation instructions - -Check `README.md` file in the root of the repository. diff --git a/advection/__init__.py b/advection/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/advection/__init__.py +++ /dev/null @@ -1,12 +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 diff --git a/advection/requirements-dev.txt b/advection/requirements-dev.txt deleted file mode 100644 index 2cb7bda5e8..0000000000 --- a/advection/requirements-dev.txt +++ /dev/null @@ -1,3 +0,0 @@ --r ../base-requirements-dev.txt --e ../common --e . diff --git a/advection/requirements.txt b/advection/requirements.txt deleted file mode 100644 index c8420e5bf8..0000000000 --- a/advection/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ --r ../base-requirements.txt -../common -. diff --git a/advection/setup.cfg b/advection/setup.cfg deleted file mode 100644 index 0d8a4b9166..0000000000 --- a/advection/setup.cfg +++ /dev/null @@ -1,50 +0,0 @@ -# This file is mainly used to configure package creation with setuptools. -# Documentation: -# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files -# -[metadata] -name = icon4py_advection -description = Icon inspired code in Python and GT4Py -long_description = file: README.md -long_description_content_type = text/markdown -url = https://github.com/C2SM/icon4py -author = ETH Zurich -author_email = gridtools@cscs.ch -license = gpl3 -license_files = LICENSE -platforms = Linux, Mac -classifiers = - Development Status :: 3 - Alpha - Intended Audience :: Science/Research - License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) - Operating System :: POSIX - Programming Language :: Python - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.10 - Programming Language :: Python :: Implementation :: CPython - Topic :: Scientific/Engineering :: Atmospheric Science - Topic :: Scientific/Engineering :: Mathematics - Topic :: Scientific/Engineering :: Physics -project_urls = - Source Code = https://github.com/C2SM/icon4py - -[options] -packages = find_namespace: -install_requires = - icon4py-common -python_requires = >=3.10 -package_dir = - = src -zip_safe = False - -[options.package_data] -# References: -# https://setuptools.pypa.io/en/latest/userguide/datafiles.html -# https://github.com/abravalheri/experiment-setuptools-package-data -* = *.md, *.rst, *.toml, *.txt, py.typed - -[options.packages.find] -where = src -exclude = - tests diff --git a/model/README.md b/model/README.md index d591915055..3f120fbe13 100644 --- a/model/README.md +++ b/model/README.md @@ -5,6 +5,7 @@ This folder contains Python implementations for multiple ICON components. It includes the following packages: - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model +- `atmosphere/advection`: Contains implementations of the advection component of the ICON model - `common`: Contains shared functionality that is required by multiple components. ## Installation Instructions @@ -16,6 +17,7 @@ In the following example it is assumed that you have already created and activat ```bash # changing into the corresponding directory cd model/atmosphere/dycore +cd model/atmosphere/advection # installing a development version pip install -r requirements-dev.txt diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py index 15dfdb0098..dab7089554 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py @@ -10,3 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py index 472c90f1f4..17eea2c697 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, where -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py index cd734f47f3..1fa9524ef0 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, sqrt, where -from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim @field_operator @@ -26,7 +26,7 @@ def _btraj_dreg_stencil_02( ) -> Field[[EdgeDim, KDim], int32]: lvn_pos = where(p_vn >= 0.0, True, False) - traj_length = sqrt(p_vn**2 + p_vt**2) * p_dt + traj_length = sqrt(p_vn*p_vn + p_vt*p_vt) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) opt_famask_dsl = where( traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 1df7ea7daa..004d758a6c 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index d08b0d0bfa..acccfd609b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim sys.setrecursionlimit(5500) @@ -32,7 +32,7 @@ def ccw( p1_lat: Field[[EdgeDim, KDim], float], p2_lon: Field[[EdgeDim, KDim], float], p2_lat: Field[[EdgeDim, KDim], float], -) -> Field[[EdgeDim, KDim], int]: +) -> Field[[EdgeDim, KDim], int32]: dx1 = p1_lon - p0_lon dy1 = p1_lat - p0_lat diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py index dde713b651..4cd783407b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.common.dimension import E2EC, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim sys.setrecursionlimit(5500) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py index 2e4a651ebc..5ffaabbe96 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py index 01fb8d4d75..1fc1cf2352 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py index 4d03e88a88..59a7281840 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py index d280b24ec9..bab061cc4e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py index c34ad0342b..54821c903a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py index b46e65be44..df6d76839a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py index 342b69e656..1c9947e4bb 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py index 474ea49602..fd8ef19b29 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py index 97dd9ead47..f912fc1f43 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, abs -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py index 395ce606b9..50a754a6ae 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py @@ -20,7 +20,7 @@ neighbor_sum, ) -from icon4py.common.dimension import C2E, C2CE, CEDim, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2E, C2CE, CEDim, C2EDim, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py index cf277deafe..6a6a9f47b4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, maximum, minimum, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py index 3abad684d3..2dea287c56 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import max_over, maximum, min_over, minimum -from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim +from icon4py.model.common.dimension import C2E2C, C2E2CDim, CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py index 55acb170e7..76e196572c 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import minimum, where -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py index 2e8e82c4f9..d645251f62 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import minimum, where -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py index e2a626bf95..47bd532a4e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum -from icon4py.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index 2777f0e036..9d1be46771 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py index 1018529602..b6ca46cbe8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import ( +from icon4py.model.common.dimension import ( C2CE, C2E, C2EDim, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py index 213d84ea1d..7fec10aad9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -267,16 +267,16 @@ def _prep_gauss_quadrature_c_list_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_y ) p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_6 = ( wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y @@ -285,10 +285,10 @@ def _prep_gauss_quadrature_c_list_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**3 - + wgt_t_detjac_2 * z_gauss_pts_2_x**3 - + wgt_t_detjac_3 * z_gauss_pts_3_x**3 - + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_8 = ( wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y @@ -297,16 +297,16 @@ def _prep_gauss_quadrature_c_list_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py index 6aa36aab23..228ed10d93 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, abs, maximum, where -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator @@ -231,16 +231,16 @@ def _prep_gauss_quadrature_c_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_y ) p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_6 = ( wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y @@ -249,28 +249,28 @@ def _prep_gauss_quadrature_c_stencil( + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**3 - + wgt_t_detjac_2 * z_gauss_pts_2_x**3 - + wgt_t_detjac_3 * z_gauss_pts_3_x**3 - + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**3 - + wgt_t_detjac_2 * z_gauss_pts_2_y**3 - + wgt_t_detjac_3 * z_gauss_pts_3_y**3 - + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y ) z_area = p_quad_vector_sum_1 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py index 4a24d8e8df..fa975ffcda 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py index fdfb97fb8e..292ce7b3a5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py @@ -17,7 +17,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast -from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim sys.setrecursionlimit(6000) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py index 03e5538bc0..b16a6d96ab 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -14,7 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py index 2e574afd0c..d3c36a5aa6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py @@ -14,7 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim +from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py index 157c5c736e..e0860c9fb3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast -from icon4py.common.dimension import CellDim +from icon4py.model.common.dimension import CellDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py index 9f3a20c26e..f9db607eaf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py index dcc530b62e..0ac10cc6a8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py index b3a029c6fe..af201f9550 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import maximum -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py index 551c28461f..85b8ac433a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, maximum -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py index 70aa04f044..a274a355b2 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py index a655fee634..804493ebc1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py index 899000f283..d3f28707fe 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index 55fd2fb9b5..7024073b68 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, where, neighbor_sum, int32, broadcast -from icon4py.common.dimension import CellDim, KDim, C2EDim, EdgeDim, C2E +from icon4py.model.common.dimension import CellDim, KDim, C2EDim, EdgeDim, C2E @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py index 71bbae9218..3e816594d8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, float64 -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py index 4bb18f5302..5cd275bfae 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32 -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py index e73f49649c..8ecbb1cb1a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py index 53b553dfcb..953b507708 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py index 3bd9bf5c1b..3ded4ed63c 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, FieldOffset, abs, int32, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py index 7315b58d3d..cba0459186 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, FieldOffset, int32, minimum, where -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 7b4cb30aa6..80d7f9d0e2 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator diff --git a/advection/setup.py b/model/atmosphere/advection/tests/conftest.py similarity index 54% rename from advection/setup.py rename to model/atmosphere/advection/tests/conftest.py index 9c9f7b81c8..7ee1151470 100644 --- a/advection/setup.py +++ b/model/atmosphere/advection/tests/conftest.py @@ -10,9 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pytest +from gt4py.next.program_processors.runners.roundtrip import executor -from setuptools import setup +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -if __name__ == "__main__": - setup() +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py index 56ce5b035f..c9a59efb6e 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import _shape, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def btraj_dreg_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py index a7f2385898..7e094d1fe2 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py @@ -14,11 +14,11 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 -from icon4py.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def btraj_dreg_stencil_02_numpy( @@ -26,7 +26,7 @@ def btraj_dreg_stencil_02_numpy( p_vt: np.array, edge_cell_length: np.array, p_dt: float, -) -> np.array: +): lvn_pos = np.where(p_vn >= 0.0, True, False) traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt @@ -35,9 +35,9 @@ def btraj_dreg_stencil_02_numpy( e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) opt_famask_dsl = np.where( - traj_length > (0.25 * np.broadcast_to(e2c_length, p_vn.shape)), - np.int32(1), - np.int32(0) + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), + int32(1), + int32(0) ) return opt_famask_dsl @@ -71,4 +71,4 @@ def test_btraj_dreg_stencil_02(): }, ) - assert np.allclose(ref, np.asarray(opt_famask_dsl)) + assert np.allclose(ref, opt_famask_dsl) diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py index 4e6e545604..c55ae59395 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py @@ -16,11 +16,11 @@ from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 -from icon4py.common.dimension import KDim, E2CDim, ECDim, EdgeDim, CellDim +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 +from icon4py.model.common.dimension import KDim, E2CDim, ECDim, EdgeDim, CellDim -from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def btraj_dreg_stencil_03_numpy( diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py index af22e53458..5c91af2e45 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py @@ -13,15 +13,15 @@ import numpy as np -from icon4py.advection.divide_flux_area_list_stencil_01 import ( +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( divide_flux_area_list_stencil_01, ) from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from gt4py.next.ffront.fbuiltins import int32 -from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim -from .test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh # FUNCTIONS # Checking turn when travelling along three points, used to check whether lines inters. @@ -44,7 +44,7 @@ def ccw( dy1dx2 = dy1 * dx2 lccw = np.where(dx1dy2 > dy1dx2, True, False) - ccw_out = np.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise return ccw_out # Checks whether two lines intersect diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py index 414963e5c0..45916e8460 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py @@ -13,15 +13,15 @@ import numpy as np -from icon4py.advection.divide_flux_area_list_stencil_02 import ( +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( divide_flux_area_list_stencil_02, ) from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from gt4py.next.ffront.fbuiltins import int32 -from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim -from .test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def divide_flux_area_list_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py index 610a3b00d0..be97aa4b4f 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import _shape, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py index 7062097cc7..5eed1821d8 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import _shape, random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py index a2735b03cc..fa827a12b3 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_02a_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py index d8f5e3c1ee..7a87257d42 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_02b_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py index 9cc38cc5f5..f4787f290a 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_02c_numpy( diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py index 2dc9e1ace0..d83b6389c0 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def face_val_ppm_stencil_05_numpy( diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py index 088759ba30..846258a19f 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.hflux_ffsl_hybrid_stencil_01a import hflux_ffsl_hybrid_stencil_01a -from icon4py.common.dimension import CellDim, KDim, EdgeDim +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import hflux_ffsl_hybrid_stencil_01a +from icon4py.model.common.dimension import CellDim, KDim, EdgeDim -from .test_utils.helpers import _shape, random_field, zero_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflux_ffsl_hybrid_stencil_01a_numpy( diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py index 065263b361..dcde5d491a 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.hflux_ffsl_hybrid_stencil_02 import hflux_ffsl_hybrid_stencil_02 -from icon4py.common.dimension import KDim, EdgeDim +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import hflux_ffsl_hybrid_stencil_02 +from icon4py.model.common.dimension import KDim, EdgeDim -from .test_utils.helpers import _shape, random_field, zero_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflux_ffsl_hybrid_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py index 0b5a4c546d..9da194cc6f 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.advection.hflx_limiter_mo_stencil_01a import hflx_limiter_mo_stencil_01a -from icon4py.common.dimension import KDim, EdgeDim, CellDim +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import hflx_limiter_mo_stencil_01a +from icon4py.model.common.dimension import KDim, EdgeDim, CellDim -from .test_utils.helpers import _shape, random_field, zero_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_01a_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py index 287aa773d9..f4088296f3 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py @@ -16,11 +16,11 @@ from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.hflx_limiter_mo_stencil_01b import hflx_limiter_mo_stencil_01b -from icon4py.common.dimension import C2EDim, CEDim, KDim, EdgeDim, CellDim +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import hflx_limiter_mo_stencil_01b +from icon4py.model.common.dimension import C2EDim, CEDim, KDim, EdgeDim, CellDim -from .test_utils.helpers import _shape, random_field, zero_field, constant_field, as_1D_sparse_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field, as_1D_sparse_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_01b_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py index 62d2339138..66cbb5212f 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -14,13 +14,13 @@ import numpy as np from numpy import int32 -from icon4py.advection.hflx_limiter_mo_stencil_02 import ( +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import ( hflx_limiter_mo_stencil_02, ) -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import constant_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py index b1935fca28..5264b1e4b4 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -12,14 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from icon4py.advection.hflx_limiter_mo_stencil_03 import ( +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_03 import ( hflx_limiter_mo_stencil_03, hflx_limiter_mo_stencil_03_min_max, ) -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_03_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py index 20b0008aa6..bd04b9b914 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -13,13 +13,13 @@ import numpy as np -from icon4py.advection.hflx_limiter_mo_stencil_04 import ( +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import ( hflx_limiter_mo_stencil_04, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_mo_stencil_04_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py index df8bbfa291..19404c50e0 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.hflx_limiter_pd_stencil_01 import ( +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import ( hflx_limiter_pd_stencil_01, ) -from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_pd_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py index 550133e0e8..dcef77e331 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -13,13 +13,13 @@ import numpy as np -from icon4py.advection.hflx_limiter_pd_stencil_02 import ( +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import ( hflx_limiter_pd_stencil_02, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import constant_field, random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import constant_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hflx_limiter_pd_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py index e009c9a04d..4a04d00962 100644 --- a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py @@ -14,11 +14,11 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.hor_adv_stencil_01 import hor_adv_stencil_01 -from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def hor_adv_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py index 4783ecaf59..12c75d4184 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -15,13 +15,13 @@ from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from gt4py.next.ffront.fbuiltins import int32 -from icon4py.advection.prep_gauss_quadrature_c_list_stencil import ( +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( prep_gauss_quadrature_c_list_stencil, ) -from icon4py.common.dimension import EdgeDim, CellDim, KDim +from icon4py.model.common.dimension import EdgeDim, CellDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def prep_gauss_quadrature_c_list_stencil_numpy( @@ -261,16 +261,16 @@ def prep_gauss_quadrature_c_list_stencil_numpy( + wgt_t_detjac_4 * z_gauss_pts_4_y ) p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_6 = ( wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y @@ -279,10 +279,10 @@ def prep_gauss_quadrature_c_list_stencil_numpy( + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**3 - + wgt_t_detjac_2 * z_gauss_pts_2_x**3 - + wgt_t_detjac_3 * z_gauss_pts_3_x**3 - + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_8 = ( wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y @@ -291,16 +291,16 @@ def prep_gauss_quadrature_c_list_stencil_numpy( + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py index 9fa2eb5e43..032041b648 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.prep_gauss_quadrature_c_stencil import ( +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( prep_gauss_quadrature_c_stencil, ) -from icon4py.common.dimension import EdgeDim, CellDim, KDim +from icon4py.model.common.dimension import EdgeDim, CellDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def prep_gauss_quadrature_c_stencil_numpy( @@ -224,16 +224,16 @@ def prep_gauss_quadrature_c_stencil_numpy( + wgt_t_detjac_4 * z_gauss_pts_4_y ) p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_6 = ( wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y @@ -242,28 +242,28 @@ def prep_gauss_quadrature_c_stencil_numpy( + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**3 - + wgt_t_detjac_2 * z_gauss_pts_2_x**3 - + wgt_t_detjac_3 * z_gauss_pts_3_x**3 - + wgt_t_detjac_4 * z_gauss_pts_4_x**3 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x ) p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y**3 - + wgt_t_detjac_2 * z_gauss_pts_2_y**3 - + wgt_t_detjac_3 * z_gauss_pts_3_y**3 - + wgt_t_detjac_4 * z_gauss_pts_4_y**3 + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y ) p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x**2 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x**2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x**2 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x**2 * z_gauss_pts_4_y + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y ) p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y**2 - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y**2 - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y**2 - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y**2 + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y ) z_area = p_quad_vector_sum_1 diff --git a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py index f998afba67..3ac9f1bd8d 100644 --- a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 -from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim +from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 +from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def rbf_intp_edge_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py index 2e45d69ead..af750be274 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.recon_lsq_cell_c_stencil import ( +from icon4py.model.atmosphere.advection.recon_lsq_cell_c_stencil import ( recon_lsq_cell_c_stencil, ) -from icon4py.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def recon_lsq_cell_c_stencil_numpy( diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py index 6246fd4520..95d0116da1 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.recon_lsq_cell_c_svd_stencil import ( +from icon4py.model.atmosphere.advection.recon_lsq_cell_c_svd_stencil import ( recon_lsq_cell_c_svd_stencil, ) -from icon4py.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def recon_lsq_cell_c_svd_stencil_numpy( c2e2c2e2c: np.ndarray, diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py index 71b8aeff0f..e31a49d40d 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.recon_lsq_cell_l_svd_stencil import ( +from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( recon_lsq_cell_l_svd_stencil, ) -from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def recon_lsq_cell_l_svd_stencil_numpy( diff --git a/model/atmosphere/advection/tests/test_set_zero_c.py b/model/atmosphere/advection/tests/test_set_zero_c.py index 89ae27a48e..059e47adcd 100644 --- a/model/atmosphere/advection/tests/test_set_zero_c.py +++ b/model/atmosphere/advection/tests/test_set_zero_c.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.set_zero_c import set_zero_c -from icon4py.common.dimension import CellDim +from icon4py.model.atmosphere.advection.set_zero_c import set_zero_c +from icon4py.model.common.dimension import CellDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def test_set_zero_cell_k(): diff --git a/model/atmosphere/advection/tests/test_set_zero_c_k.py b/model/atmosphere/advection/tests/test_set_zero_c_k.py index 60ab010379..6f26b960fd 100644 --- a/model/atmosphere/advection/tests/test_set_zero_c_k.py +++ b/model/atmosphere/advection/tests/test_set_zero_c_k.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.set_zero_c_k import set_zero_c_k -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.set_zero_c_k import set_zero_c_k +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def test_set_zero_c_k(): diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py index 7cab0293ee..803c9e62eb 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def step_advection_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py index 030d190255..bc60330323 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.step_advection_stencil_02 import step_advection_stencil_02 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def step_advection_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py index d3aac81733..e69a82f6d0 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.step_advection_stencil_03 import step_advection_stencil_03 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def step_advection_stencil_03_numpy( diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py index 0b7178a13e..36addc80e9 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def step_advection_stencil_04_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py index f1b1519a51..3bb27e469c 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py @@ -13,14 +13,14 @@ import numpy as np -from icon4py.advection.upwind_hflux_miura3_stencil_01 import ( +from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( upwind_hflux_miura3_stencil_01, ) from gt4py.next.ffront.fbuiltins import int32 -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field, random_mask -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura3_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py index 59cf65e96b..1939b09564 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -14,11 +14,11 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.advection.upwind_hflux_miura_cycl_stencil_01 import upwind_hflux_miura_cycl_stencil_01 -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import upwind_hflux_miura_cycl_stencil_01 +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import _shape, random_field, zero_field, random_mask -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, random_mask +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_cycl_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py index ca075b3022..a89ae88c34 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.advection.upwind_hflux_miura_cycl_stencil_02 import upwind_hflux_miura_cycl_stencil_02 -from icon4py.common.dimension import CellDim, EdgeDim, KDim, CEDim, C2EDim +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import upwind_hflux_miura_cycl_stencil_02 +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, CEDim, C2EDim -from .test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_cycl_stencil_02_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py index 1c3e005af5..52110a21a5 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.upwind_hflux_miura_cycl_stencil_03a import upwind_hflux_miura_cycl_stencil_03a -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import upwind_hflux_miura_cycl_stencil_03a +from icon4py.model.common.dimension import EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_cycl_stencil_03a_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py index 2e8df339ed..a3af39a827 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.upwind_hflux_miura_cycl_stencil_03b import upwind_hflux_miura_cycl_stencil_03b -from icon4py.common.dimension import EdgeDim, KDim +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import upwind_hflux_miura_cycl_stencil_03b +from icon4py.model.common.dimension import EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_cycl_stencil_03b_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py index faf80d6030..5d65379c80 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -13,14 +13,14 @@ import numpy as np -from icon4py.advection.upwind_hflux_miura_stencil_01 import ( +from icon4py.model.atmosphere.advection.upwind_hflux_miura_stencil_01 import ( upwind_hflux_miura_stencil_01, ) from gt4py.next.ffront.fbuiltins import int32 -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field, constant_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field, constant_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_hflux_miura_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 37edce8962..6321d78c8d 100644 --- a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -13,13 +13,13 @@ import numpy as np -from icon4py.advection.upwind_vflux_ppm_stencil_01 import ( +from icon4py.model.atmosphere.advection.upwind_vflux_ppm_stencil_01 import ( upwind_vflux_ppm_stencil_01, ) -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def upwind_vflux_ppm_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py index 9b6261745f..f9b2586bfd 100644 --- a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.advection.vert_adv_stencil_01 import vert_adv_stencil_01 -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def vert_adv_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py index e0e475bc0c..163eeb900d 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.advection.v_limit_prbl_sm_stencil_01 import ( +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import ( v_limit_prbl_sm_stencil_01, ) -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def v_limit_prbl_sm_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py index f98cb6ee90..670664f1f0 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -14,13 +14,13 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.advection.v_limit_prbl_sm_stencil_02 import ( +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import ( v_limit_prbl_sm_stencil_02, ) -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim -from .test_utils.helpers import random_field, random_mask, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def v_limit_prbl_sm_stencil_02_numpy( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py similarity index 95% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py index 375755f56c..f23101a69a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/compute_airmass.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py @@ -14,7 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, KDim @field_operator diff --git a/atm_dyn_iconam/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py similarity index 100% rename from atm_dyn_iconam/tests/test_compute_airmass.py rename to model/atmosphere/dycore/tests/test_compute_airmass.py diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index 2cff62c6eb..d6f92a2cb8 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -33,6 +33,7 @@ V2CDim, V2EDim, VertexDim, + C2E2C2E2CDim, ) @@ -347,6 +348,28 @@ class SimpleMeshData: ] ) + c2e2c2e2c_table = np.asarray( + [ + [15, 4, 3,12, 14, 1, 7, 6, 2], #1c + [16, 5, 4,12, 13, 2, 8, 7, 0], + [17, 3, 5,13, 14, 0, 6, 8, 1], + [0, 6, 2, 17, 5, 9, 10, 15, 4], + [1, 7, 0, 15, 3, 16, 5, 10, 11], #5c + [2, 8, 1, 4, 16, 17, 3, 9, 11], + [3, 10, 9, 2, 0, 7, 13, 8, 12], + [4, 11, 10, 0, 1, 8, 14, 6, 13], + [5, 9, 11, 1, 2, 3, 12, 7, 14], + [6, 12, 8, 5, 11, 3, 10, 16, 15], #10c + [7, 13, 6, 3, 9, 4, 11, 16, 17], + [8, 14, 7, 4, 10, 5, 9, 15, 17], + [9, 16, 15, 8, 6, 1, 13, 0, 14], + [10, 17, 16, 6, 7, 2, 14, 1, 12], + [11, 15, 17, 7, 8, 2, 13, 0, 12], #15c + [12, 0, 14,11, 17, 9, 16, 3, 4], + [13, 1, 12, 9, 15, 10, 17, 4, 5], + [14, 2, 13,10, 16, 5, 3, 11, 15], + ] + ) class SimpleMesh: _DEFAULT_K_LEVEL = 10 @@ -363,6 +386,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.e2c2v = SimpleMeshData.e2c2v_table self.v2c = SimpleMeshData.v2c_table self.v2e = SimpleMeshData.v2e_table + self.c2e2c2e2c = SimpleMeshData.c2e2c2e2c_table self.n_e2c = self.e2c.shape[1] self.n_e2v = self.e2v.shape[1] self.n_c2e = self.c2e.shape[1] @@ -374,6 +398,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_v2c = self.v2c.shape[1] self.n_v2e = self.v2e.shape[1] self.n_cells = self.c2e.shape[0] + self.n_c2e2c2e2c = self.c2e2c2e2c.shape[1] self.n_edges = 27 self.n_vertices = 9 self.k_level = k_level @@ -388,12 +413,14 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): E2C2EODim: self.n_e2c2eO, E2C2EDim: self.n_e2c2e, V2CDim: self.n_v2c, + V2CDim: self.n_v2c, KDim: self.k_level, VertexDim: self.n_vertices, V2EDim: self.n_v2e, CEDim: self.n_cells * self.n_c2e, E2C2VDim: self.n_e2c2v, ECVDim: self.n_edges * self.n_e2c2v, + C2E2C2E2CDim: self.n_c2e2c2e2c, } def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: @@ -426,6 +453,9 @@ def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) + def get_c2e2c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e2c2e2c, CellDim, CellDim, self.n_c2e2c2e2c) + def get_offset_provider(self): return { "C2E": self.get_c2e_offset_provider(), @@ -439,4 +469,5 @@ def get_offset_provider(self): "E2V": self.get_e2v_offset_provider(), "E2C2V": self.get_e2c2v_offset_provider(), "Koff": KDim, + "C2E2C2E2C": self.get_c2e2c2e2c_offset_provider(), } diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index 1b7e16d5a1..10bde64b68 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -1,3 +1,4 @@ -r ../base-requirements-dev.txt -e ./atmosphere/dycore +-e ./atmosphere/advection -e ./common diff --git a/model/requirements.txt b/model/requirements.txt index 2045fe9058..4406060faf 100644 --- a/model/requirements.txt +++ b/model/requirements.txt @@ -1,3 +1,4 @@ -r ../base-requirements.txt ./atmosphere/dycore +./atmosphere/advection ./common diff --git a/model/tox.ini b/model/tox.ini index 27fb18d610..b97a24f41e 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -15,6 +15,7 @@ deps = -r {toxinidir}/requirements-dev.txt commands = -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/advection/src common/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 1b98413f77..8704e2936d 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -29,7 +29,7 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atmosphere.dycore"] + _STENCIL_PACKAGES = ["atmosphere.dycore", "atmosphere.advection"] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 898b6b71a0..a3d6e7ba57 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -16,6 +16,7 @@ import re import icon4py.model.atmosphere.dycore as dycore +import icon4py.model.atmosphere.advection as advection import pytest from click.testing import CliRunner from icon4pytools.icon4pygen.cli import main @@ -42,6 +43,11 @@ def dycore_fencils() -> list[tuple[str, str]]: fencils = [(DYCORE_PKG, stencil) for stencil in stencils] return fencils +def advection_fencils() -> list[tuple[str, str]]: + pkgpath = os.path.dirname(advection.__file__) + stencils = [name for _, name, _ in pkgutil.iter_modules([pkgpath])] + fencils = [(DYCORE_PKG, stencil) for stencil in stencils] + return fencils def check_cpp_codegen(fname: str) -> None: stencil_name = fname.replace(".cpp", "") @@ -120,6 +126,14 @@ def test_codegen_dycore(cli, stencil_module, stencil_name) -> None: assert result.exit_code == 0 check_code_was_generated(stencil_name) +@pytest.mark.parametrize(("stencil_module", "stencil_name"), advection_fencils()) +def test_codegen_advection(cli, stencil_module, stencil_name) -> None: + module_path = get_stencil_module_path(stencil_module, stencil_name) + with cli.isolated_filesystem(): + result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) + assert result.exit_code == 0 + check_code_was_generated(stencil_name) + def test_invalid_module_path(cli) -> None: module_path = get_stencil_module_path("some_module", "foo") diff --git a/tox.ini b/tox.ini index 65d9765d16..4768cc709e 100644 --- a/tox.ini +++ b/tox.ini @@ -15,6 +15,7 @@ deps = -r {toxinidir}/requirements-dev.txt commands = -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/common/src tools/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/advection/src model/common/src tools/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html From 27b89e9aadeaecb2199dbf88d9e1612ae11b07d7 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 21 Aug 2023 09:59:01 +0200 Subject: [PATCH 086/105] Pre-commit run - ICON4Py Quality Assurance --- .../advection/btraj_dreg_stencil_02.py | 2 +- .../advection/btraj_dreg_stencil_03.py | 2 + .../divide_flux_area_list_stencil_01.py | 264 +++++++++---- .../hflux_ffsl_hybrid_stencil_01a.py | 81 +++- .../advection/hflx_limiter_mo_stencil_01b.py | 35 +- .../advection/hflx_limiter_pd_stencil_01.py | 9 +- .../model/atmosphere/advection/py.typed | 0 .../advection/recon_lsq_cell_c_stencil.py | 8 +- .../advection/recon_lsq_cell_c_svd_stencil.py | 8 +- .../advection/recon_lsq_cell_l_svd_stencil.py | 5 +- .../upwind_hflux_miura3_stencil_01.py | 85 ++++- .../upwind_hflux_miura_cycl_stencil_01.py | 30 +- .../upwind_hflux_miura_cycl_stencil_02.py | 37 +- .../upwind_hflux_miura_cycl_stencil_03a.py | 8 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 11 +- .../upwind_hflux_miura_stencil_01.py | 24 +- .../tests/test_btraj_dreg_stencil_01.py | 16 +- .../tests/test_btraj_dreg_stencil_02.py | 28 +- .../tests/test_btraj_dreg_stencil_03.py | 23 +- .../test_divide_flux_area_list_stencil_01.py | 357 +++++++++++++----- .../test_divide_flux_area_list_stencil_02.py | 83 ++-- .../tests/test_face_val_ppm_stencil_01.py | 11 +- .../tests/test_face_val_ppm_stencil_02.py | 5 +- .../tests/test_face_val_ppm_stencil_02a.py | 5 +- .../tests/test_face_val_ppm_stencil_02b.py | 5 +- .../tests/test_face_val_ppm_stencil_02c.py | 5 +- .../tests/test_face_val_ppm_stencil_05.py | 5 +- .../test_hflux_ffsl_hybrid_stencil_01a.py | 105 ++++-- .../test_hflux_ffsl_hybrid_stencil_02.py | 14 +- .../tests/test_hflx_limiter_mo_stencil_01a.py | 15 +- .../tests/test_hflx_limiter_mo_stencil_01b.py | 51 +-- .../tests/test_hflx_limiter_mo_stencil_02.py | 7 +- .../tests/test_hflx_limiter_mo_stencil_03.py | 1 - .../tests/test_hflx_limiter_mo_stencil_04.py | 1 - .../tests/test_hflx_limiter_pd_stencil_01.py | 7 +- .../tests/test_hflx_limiter_pd_stencil_02.py | 1 - .../tests/test_hor_adv_stencil_01.py | 10 +- ...st_prep_gauss_quadrature_c_list_stencil.py | 139 +++---- .../test_prep_gauss_quadrature_c_stencil.py | 124 +++--- .../tests/test_rbf_intp_edge_stencil_01.py | 5 +- .../tests/test_recon_lsq_cell_c_stencil.py | 163 ++++---- .../test_recon_lsq_cell_c_svd_stencil.py | 86 +++-- .../test_recon_lsq_cell_l_svd_stencil.py | 7 +- .../advection/tests/test_set_zero_c.py | 1 - .../advection/tests/test_set_zero_c_k.py | 1 - .../tests/test_step_advection_stencil_01.py | 5 +- .../tests/test_step_advection_stencil_02.py | 5 +- .../tests/test_step_advection_stencil_03.py | 5 +- .../tests/test_step_advection_stencil_04.py | 5 +- .../test_upwind_hflux_miura3_stencil_01.py | 87 ++++- ...test_upwind_hflux_miura_cycl_stencil_01.py | 35 +- ...test_upwind_hflux_miura_cycl_stencil_02.py | 31 +- ...est_upwind_hflux_miura_cycl_stencil_03a.py | 8 +- ...est_upwind_hflux_miura_cycl_stencil_03b.py | 12 +- .../test_upwind_hflux_miura_stencil_01.py | 30 +- .../tests/test_upwind_vflux_ppm_stencil_01.py | 1 - .../tests/test_vert_adv_stencil_01.py | 5 +- .../tests/test_vlimit_prbl_sm_stencil_01.py | 1 - .../tests/test_vlimit_prbl_sm_stencil_02.py | 7 +- .../dycore/apply_diffusion_to_vn.py | 4 +- ...ute_horizontal_gradients_for_turbulance.py | 8 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 4 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 4 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 4 +- .../calculate_diagnostics_for_turbulence.py | 4 +- ..._coefficients_for_grid_point_cold_pools.py | 4 +- ...ate_horizontal_gradients_for_turbulence.py | 4 +- .../dycore/calculate_nabla2_for_theta.py | 4 +- ...n_coefficient_for_grid_point_cold_pools.py | 4 +- .../mo_advection_traj_btraj_compute_o1_dsl.py | 18 +- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 4 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_02.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_10.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 +- .../dycore/mo_solve_nonhydro_stencil_17.py | 5 +- .../dycore/mo_solve_nonhydro_stencil_18.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_20.py | 10 +- .../dycore/mo_solve_nonhydro_stencil_22.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_24.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_39.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_40.py | 12 +- .../dycore/mo_solve_nonhydro_stencil_42.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_55.py | 13 +- .../dycore/mo_solve_nonhydro_stencil_56_63.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_58.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_60.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_65.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_67.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_68.py | 5 +- .../mo_velocity_advection_stencil_04.py | 4 +- .../mo_velocity_advection_stencil_07.py | 4 +- .../mo_velocity_advection_stencil_10.py | 4 +- .../mo_velocity_advection_stencil_14.py | 8 +- .../mo_velocity_advection_stencil_16.py | 8 +- .../mo_velocity_advection_stencil_18.py | 3 +- .../mo_velocity_advection_stencil_19.py | 5 +- .../mo_velocity_advection_stencil_20.py | 4 +- ...fusion_nabla_of_theta_over_steep_points.py | 15 +- .../dycore/update_theta_and_exner.py | 4 +- ...st_apply_nabla2_and_nabla4_global_to_vn.py | 8 +- .../dycore/tests/test_apply_nabla2_to_w.py | 4 +- ...st_calculate_diagnostics_for_turbulence.py | 8 +- .../tests/test_calculate_nabla2_of_theta.py | 4 +- .../dycore/tests/test_compute_airmass.py | 1 - ..._mo_advection_traj_btraj_compute_o1_dsl.py | 18 +- .../test_mo_solve_nonhydro_stencil_02.py | 4 +- .../test_mo_solve_nonhydro_stencil_04.py | 9 +- .../test_mo_solve_nonhydro_stencil_06.py | 4 +- .../test_mo_solve_nonhydro_stencil_15.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 23 +- .../test_mo_solve_nonhydro_stencil_19.py | 5 +- .../test_mo_solve_nonhydro_stencil_20.py | 8 +- .../test_mo_solve_nonhydro_stencil_21.py | 4 +- .../test_mo_solve_nonhydro_stencil_25.py | 4 +- .../test_mo_solve_nonhydro_stencil_26.py | 4 +- .../test_mo_solve_nonhydro_stencil_28.py | 4 +- .../test_mo_solve_nonhydro_stencil_29.py | 4 +- .../test_mo_solve_nonhydro_stencil_42.py | 4 +- .../test_mo_solve_nonhydro_stencil_44.py | 4 +- .../test_mo_solve_nonhydro_stencil_47.py | 4 +- .../test_mo_solve_nonhydro_stencil_51.py | 4 +- .../test_mo_solve_nonhydro_stencil_52.py | 8 +- .../test_mo_solve_nonhydro_stencil_54.py | 4 +- .../test_mo_solve_nonhydro_stencil_55.py | 5 +- .../test_mo_solve_nonhydro_stencil_58.py | 4 +- .../test_mo_solve_nonhydro_stencil_60.py | 4 +- .../test_mo_solve_nonhydro_stencil_62.py | 4 +- .../test_mo_solve_nonhydro_stencil_65.py | 4 +- .../test_mo_solve_nonhydro_stencil_68.py | 5 +- .../test_mo_velocity_advection_stencil_02.py | 8 +- .../test_mo_velocity_advection_stencil_03.py | 8 +- .../test_mo_velocity_advection_stencil_06.py | 8 +- .../test_mo_velocity_advection_stencil_08.py | 4 +- .../test_mo_velocity_advection_stencil_09.py | 4 +- .../test_mo_velocity_advection_stencil_10.py | 4 +- .../test_mo_velocity_advection_stencil_13.py | 4 +- .../test_mo_velocity_advection_stencil_19.py | 5 +- ...fusion_nabla_of_theta_over_steep_points.py | 7 +- 139 files changed, 1756 insertions(+), 825 deletions(-) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py index 1fa9524ef0..5984b77c5a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -26,7 +26,7 @@ def _btraj_dreg_stencil_02( ) -> Field[[EdgeDim, KDim], int32]: lvn_pos = where(p_vn >= 0.0, True, False) - traj_length = sqrt(p_vn*p_vn + p_vt*p_vt) * p_dt + traj_length = sqrt(p_vn * p_vn + p_vt * p_vt) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) opt_famask_dsl = where( traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 004d758a6c..bc120fe1b9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -1,3 +1,5 @@ +# TODO: This license is not consistent with license used in the project. +# Delete the inconsistent license and above line and rerun pre-commit to insert a good license. # ICON4Py - ICON inspired code in Python and GT4Py # # Copyright (c) 2022, ETH Zurich and MeteoSwiss diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index acccfd609b..f240bc0ccf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -258,76 +258,136 @@ def _divide_flux_area_list_stencil_01( ) # Case 1 - patch 0 - dreg_patch0_1_lon_dsl = where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = where( + mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = where( + mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = where( - mask_case1, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + mask_case1, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = where( - mask_case1, where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + mask_case1, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, ) dreg_patch0_3_lon_dsl = where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) dreg_patch0_3_lat_dsl = where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = where( - mask_case1, where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + mask_case1, + where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) dreg_patch0_4_lat_dsl = where( - mask_case1, where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + mask_case1, + where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 1 - patch 1 dreg_patch1_1_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) dreg_patch1_1_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) dreg_patch1_4_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) dreg_patch1_4_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = where(mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0) - dreg_patch1_2_lat_vmask = where(mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0) - dreg_patch1_3_lon_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0) - dreg_patch1_3_lat_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0) + dreg_patch1_2_lon_vmask = where( + mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = where( + mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) # Case 1 - patch 2 dreg_patch2_1_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) dreg_patch2_1_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) dreg_patch2_4_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) dreg_patch2_4_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0) - dreg_patch2_2_lat_vmask = where(mask_case1, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0) - dreg_patch2_3_lon_vmask = where(mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0) - dreg_patch2_3_lat_vmask = where(mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0) + dreg_patch2_2_lon_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = where( + mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = where( + mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) # ------------------------------------------------- Case 2a mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool # Case 2a - patch 0 - dreg_patch0_1_lon_dsl = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = where( - mask_case2a, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + mask_case2a, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = where( - mask_case2a, where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + mask_case2a, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where( + mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl + ) + dreg_patch0_3_lat_dsl = where( + mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl ) - dreg_patch0_3_lon_dsl = where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = where( - mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + mask_case2a, + where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) dreg_patch0_4_lat_dsl = where( - mask_case2a, where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + mask_case2a, + where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 2a - patch 1 - dreg_patch1_1_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) - dreg_patch1_4_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_4_lon_vmask = where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask + ) + dreg_patch1_4_lat_vmask = where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask + ) dreg_patch1_2_lon_vmask = where( - mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask + mask_case2a, + where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, ) dreg_patch1_2_lat_vmask = where( - mask_case2a, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask + mask_case2a, + where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, ) dreg_patch1_3_lon_vmask = where( - mask_case2a, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask + mask_case2a, + where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, ) dreg_patch1_3_lat_vmask = where( - mask_case2a, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask + mask_case2a, + where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, ) # Case 2a - patch 2 dreg_patch2_1_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) @@ -342,8 +402,12 @@ def _divide_flux_area_list_stencil_01( # -------------------------------------------------- Case 2b mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool # Case 2b - patch 0 - dreg_patch0_1_lon_dsl = where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = where( + mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = where( + mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = where( mask_case2b, where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -376,21 +440,37 @@ def _divide_flux_area_list_stencil_01( dreg_patch1_4_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) # Case 2b - patch 2 - dreg_patch2_1_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) - dreg_patch2_4_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_4_lon_vmask = where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask + ) + dreg_patch2_4_lat_vmask = where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask + ) dreg_patch2_2_lon_vmask = where( - mask_case2b, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask + mask_case2b, + where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, ) dreg_patch2_2_lat_vmask = where( - mask_case2b, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask + mask_case2b, + where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, ) dreg_patch2_3_lon_vmask = where( - mask_case2b, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask + mask_case2b, + where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, ) dreg_patch2_3_lat_vmask = where( - mask_case2b, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask + mask_case2b, + where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, ) # flux area edge 1 and 2 @@ -427,8 +507,12 @@ def _divide_flux_area_list_stencil_01( tri_line1_p2_lat, ) # Case 3a - patch 0 - dreg_patch0_1_lon_dsl = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = where( mask_case3a, where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -452,21 +536,37 @@ def _divide_flux_area_list_stencil_01( dreg_patch0_4_lat_dsl, ) # Case 3a - patch 1 - dreg_patch1_1_lon_vmask = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_1_lon_vmask = where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) dreg_patch1_2_lon_vmask = where( - mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask + mask_case3a, + where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, ) dreg_patch1_2_lat_vmask = where( - mask_case3a, where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask + mask_case3a, + where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = where( + mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask + ) + dreg_patch1_3_lat_vmask = where( + mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask ) - dreg_patch1_3_lon_vmask = where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) dreg_patch1_4_lon_vmask = where( - mask_case3a, where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, ) dreg_patch1_4_lat_vmask = where( - mask_case3a, where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, ) # Case 3a - patch 2 dreg_patch2_1_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) @@ -502,21 +602,37 @@ def _divide_flux_area_list_stencil_01( tri_line2_p2_lat, ) # Case 3b - patch 0 - dreg_patch0_1_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_1_lon_dsl = where( + mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = where( + mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) + dreg_patch0_4_lon_dsl = where( + mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl + ) + dreg_patch0_4_lat_dsl = where( + mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl + ) dreg_patch0_2_lon_dsl = where( - mask_case3b, where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl + mask_case3b, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = where( - mask_case3b, where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl + mask_case3b, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, ) dreg_patch0_3_lon_dsl = where( - mask_case3b, where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl + mask_case3b, + where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, ) dreg_patch0_3_lat_dsl = where( - mask_case3b, where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl + mask_case3b, + where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, ) # Case 3b - patch 1 dreg_patch1_1_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) @@ -528,21 +644,37 @@ def _divide_flux_area_list_stencil_01( dreg_patch1_4_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 - dreg_patch2_1_lon_vmask = where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_1_lon_vmask = where( + mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = where( + mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) dreg_patch2_2_lon_vmask = where( - mask_case3b, where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask + mask_case3b, + where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, ) dreg_patch2_2_lat_vmask = where( - mask_case3b, where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask + mask_case3b, + where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = where( + mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask + ) + dreg_patch2_3_lat_vmask = where( + mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask ) - dreg_patch2_3_lon_vmask = where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) dreg_patch2_4_lon_vmask = where( - mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask + mask_case3b, + where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, ) dreg_patch2_4_lat_vmask = where( - mask_case3b, where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask + mask_case3b, + where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, ) # --------------------------------------------- Case 4 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py index 1c9947e4bb..7457dc3ebc 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -43,20 +43,71 @@ def _hflux_ffsl_hybrid_stencil_01a( patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e_hybrid_1a = ( - where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) * z_quad_vector_sum0_1 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) * z_quad_vector_sum0_2 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) * z_quad_vector_sum0_3 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_4(E2C[1]), z_lsq_coeff_4(E2C[0])) * z_quad_vector_sum0_4 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_5(E2C[1]), z_lsq_coeff_5(E2C[0])) * z_quad_vector_sum0_5 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_6(E2C[1]), z_lsq_coeff_6(E2C[0])) * z_quad_vector_sum0_6 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_7(E2C[1]), z_lsq_coeff_7(E2C[0])) * z_quad_vector_sum0_7 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_8(E2C[1]), z_lsq_coeff_8(E2C[0])) * z_quad_vector_sum0_8 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_9(E2C[1]), z_lsq_coeff_9(E2C[0])) * z_quad_vector_sum0_9 - + where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0])) * z_quad_vector_sum0_10 - ) + p_out_e_hybrid_1a = ( + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1(E2C[1]), + z_lsq_coeff_1(E2C[0]), + ) + * z_quad_vector_sum0_1 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2(E2C[1]), + z_lsq_coeff_2(E2C[0]), + ) + * z_quad_vector_sum0_2 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3(E2C[1]), + z_lsq_coeff_3(E2C[0]), + ) + * z_quad_vector_sum0_3 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4(E2C[1]), + z_lsq_coeff_4(E2C[0]), + ) + * z_quad_vector_sum0_4 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5(E2C[1]), + z_lsq_coeff_5(E2C[0]), + ) + * z_quad_vector_sum0_5 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6(E2C[1]), + z_lsq_coeff_6(E2C[0]), + ) + * z_quad_vector_sum0_6 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7(E2C[1]), + z_lsq_coeff_7(E2C[0]), + ) + * z_quad_vector_sum0_7 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8(E2C[1]), + z_lsq_coeff_8(E2C[0]), + ) + * z_quad_vector_sum0_8 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9(E2C[1]), + z_lsq_coeff_9(E2C[0]), + ) + * z_quad_vector_sum0_9 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10(E2C[1]), + z_lsq_coeff_10(E2C[0]), + ) + * z_quad_vector_sum0_10 + ) + + return p_out_e_hybrid_1a - return p_out_e_hybrid_1a @program def hflux_ffsl_hybrid_stencil_01a( @@ -105,5 +156,5 @@ def hflux_ffsl_hybrid_stencil_01a( z_quad_vector_sum0_9, z_quad_vector_sum0_10, patch0_cell_rel_idx_dsl, - out=(p_out_e_hybrid_1a) + out=(p_out_e_hybrid_1a), ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py index 50a754a6ae..fcc5e26405 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py @@ -20,7 +20,15 @@ neighbor_sum, ) -from icon4py.model.common.dimension import C2E, C2CE, CEDim, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import ( + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) @field_operator @@ -41,20 +49,21 @@ def _hflx_limiter_mo_stencil_01b( ]: zero = broadcast(0.0, (CellDim, KDim)) - z_mflx_anti_1 = ( p_dtime * geofac_div(C2CE[0]) / p_rhodz_new - * z_anti(C2E[0]) ) - z_mflx_anti_2 = ( p_dtime * geofac_div(C2CE[1]) / p_rhodz_new - * z_anti(C2E[1]) ) - z_mflx_anti_3 = ( p_dtime * geofac_div(C2CE[2]) / p_rhodz_new - * z_anti(C2E[2]) ) + z_mflx_anti_1 = p_dtime * geofac_div(C2CE[0]) / p_rhodz_new * z_anti(C2E[0]) + z_mflx_anti_2 = p_dtime * geofac_div(C2CE[1]) / p_rhodz_new * z_anti(C2E[1]) + z_mflx_anti_3 = p_dtime * geofac_div(C2CE[2]) / p_rhodz_new * z_anti(C2E[2]) - z_mflx_anti_in = -1.0 * ( minimum(zero, z_mflx_anti_1) - + minimum(zero, z_mflx_anti_2) - + minimum(zero, z_mflx_anti_3) ) + z_mflx_anti_in = -1.0 * ( + minimum(zero, z_mflx_anti_1) + + minimum(zero, z_mflx_anti_2) + + minimum(zero, z_mflx_anti_3) + ) - z_mflx_anti_out = ( maximum(zero, z_mflx_anti_1) - + maximum(zero, z_mflx_anti_2) - + maximum(zero, z_mflx_anti_3) ) + z_mflx_anti_out = ( + maximum(zero, z_mflx_anti_1) + + maximum(zero, z_mflx_anti_2) + + maximum(zero, z_mflx_anti_3) + ) z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div(C2CE), axis=C2EDim) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py index 47bd532a4e..cc51b56438 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py @@ -14,7 +14,14 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum -from icon4py.model.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import ( + C2CE, + C2E, + CEDim, + CellDim, + EdgeDim, + KDim, +) @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py index 292ce7b3a5..d22089f684 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py @@ -17,7 +17,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast -from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import ( + C2CECEC, + C2E2C2E2C, + CECECDim, + CellDim, + KDim, +) sys.setrecursionlimit(6000) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py index b16a6d96ab..edffe26fcb 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -14,7 +14,13 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim +from icon4py.model.common.dimension import ( + C2CECEC, + C2E2C2E2C, + CECECDim, + CellDim, + KDim, +) @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py index d3c36a5aa6..5e85447556 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py @@ -51,5 +51,8 @@ def recon_lsq_cell_l_svd_stencil( p_coeff_3_dsl: Field[[CellDim, KDim], float], ): _recon_lsq_cell_l_svd_stencil( - p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl) + p_cc, + lsq_pseudoinv_1, + lsq_pseudoinv_2, + out=(p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl), ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py index 804493ebc1..50a1be5a90 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -45,20 +45,75 @@ def _upwind_hflux_miura3_stencil_01( cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e_miura3 = ( - where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) * z_quad_vector_sum_1 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) * z_quad_vector_sum_2 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) * z_quad_vector_sum_3 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_4(E2C[1]), z_lsq_coeff_4(E2C[0])) * z_quad_vector_sum_4 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_5(E2C[1]), z_lsq_coeff_5(E2C[0])) * z_quad_vector_sum_5 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_6(E2C[1]), z_lsq_coeff_6(E2C[0])) * z_quad_vector_sum_6 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_7(E2C[1]), z_lsq_coeff_7(E2C[0])) * z_quad_vector_sum_7 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_8(E2C[1]), z_lsq_coeff_8(E2C[0])) * z_quad_vector_sum_8 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_9(E2C[1]), z_lsq_coeff_9(E2C[0])) * z_quad_vector_sum_9 - + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0])) * z_quad_vector_sum_10 - ) / z_dreg_area * p_mass_flx_e + p_out_e_miura3 = ( + ( + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1(E2C[1]), + z_lsq_coeff_1(E2C[0]), + ) + * z_quad_vector_sum_1 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2(E2C[1]), + z_lsq_coeff_2(E2C[0]), + ) + * z_quad_vector_sum_2 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3(E2C[1]), + z_lsq_coeff_3(E2C[0]), + ) + * z_quad_vector_sum_3 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4(E2C[1]), + z_lsq_coeff_4(E2C[0]), + ) + * z_quad_vector_sum_4 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5(E2C[1]), + z_lsq_coeff_5(E2C[0]), + ) + * z_quad_vector_sum_5 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6(E2C[1]), + z_lsq_coeff_6(E2C[0]), + ) + * z_quad_vector_sum_6 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7(E2C[1]), + z_lsq_coeff_7(E2C[0]), + ) + * z_quad_vector_sum_7 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8(E2C[1]), + z_lsq_coeff_8(E2C[0]), + ) + * z_quad_vector_sum_8 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9(E2C[1]), + z_lsq_coeff_9(E2C[0]), + ) + * z_quad_vector_sum_9 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10(E2C[1]), + z_lsq_coeff_10(E2C[0]), + ) + * z_quad_vector_sum_10 + ) + / z_dreg_area + * p_mass_flx_e + ) + + return p_out_e_miura3 - return p_out_e_miura3 @program def upwind_hflux_miura3_stencil_01( @@ -111,5 +166,5 @@ def upwind_hflux_miura3_stencil_01( z_dreg_area, p_mass_flx_e, cell_rel_idx_dsl, - out=(p_out_e_miura3) + out=(p_out_e_miura3), ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py index d3f28707fe..eef9ba61e9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -29,12 +29,28 @@ def _upwind_hflux_miura_cycl_stencil_01( cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - z_tracer_mflx_dsl = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl(E2C[1]), z_lsq_coeff_1_dsl(E2C[0])) - + distv_bary_1 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_dsl(E2C[1]), z_lsq_coeff_2_dsl(E2C[0])) - + distv_bary_2 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_dsl(E2C[1]), z_lsq_coeff_3_dsl(E2C[0])) - ) * p_mass_flx_e + z_tracer_mflx_dsl = ( + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl(E2C[1]), + z_lsq_coeff_1_dsl(E2C[0]), + ) + + distv_bary_1 + * where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl(E2C[1]), + z_lsq_coeff_2_dsl(E2C[0]), + ) + + distv_bary_2 + * where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl(E2C[1]), + z_lsq_coeff_3_dsl(E2C[0]), + ) + ) * p_mass_flx_e + + return z_tracer_mflx_dsl - return z_tracer_mflx_dsl @program def upwind_hflux_miura_cycl_stencil_01( @@ -55,5 +71,5 @@ def upwind_hflux_miura_cycl_stencil_01( distv_bary_2, p_mass_flx_e, cell_rel_idx_dsl, - out=(z_tracer_mflx_dsl) + out=(z_tracer_mflx_dsl), ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index 7024073b68..cb7c1809af 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -12,9 +12,15 @@ # 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, neighbor_sum, int32, broadcast +from gt4py.next.ffront.fbuiltins import ( + Field, + broadcast, + int32, + neighbor_sum, + where, +) -from icon4py.model.common.dimension import CellDim, KDim, C2EDim, EdgeDim, C2E +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -27,19 +33,26 @@ def _upwind_hflux_miura_cycl_stencil_02( z_rho_now: Field[[CellDim, KDim], float], z_tracer_now: Field[[CellDim, KDim], float], z_dtsub: float, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + + z_rhofluxdiv_c_out = ( + neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) + if nsub == int32(1) + else z_rhofluxdiv_c + ) - z_rhofluxdiv_c_out = neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) if nsub == int32(1) else z_rhofluxdiv_c - z_fluxdiv_c_dsl = neighbor_sum(z_tracer_mflx(C2E) * geofac_div, axis=C2EDim) - - z_rho_new_dsl = ( z_rho_now - - z_dtsub * z_rhofluxdiv_c_out ) - z_tracer_new_dsl = (( z_tracer_now - * z_rho_now - - z_dtsub * z_fluxdiv_c_dsl ) - / z_rho_new_dsl) + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out + + z_tracer_new_dsl = ( + z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl + ) / z_rho_new_dsl return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py index 3e816594d8..7c27ef4ee9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -13,7 +13,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import int32, float64 +from gt4py.next.ffront.fbuiltins import float64, int32 from icon4py.model.common.dimension import EdgeDim, KDim @@ -23,7 +23,7 @@ def _upwind_hflux_miura_cycl_stencil_03a( z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl)/float64(2) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float64(2) return p_out_e @@ -34,7 +34,5 @@ def upwind_hflux_miura_cycl_stencil_03a( p_out_e: Field[[EdgeDim, KDim], float], ): _upwind_hflux_miura_cycl_stencil_03a( - z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl, - out=(p_out_e) + z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e) ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py index 5cd275bfae..d127c50224 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -24,9 +24,9 @@ def _upwind_hflux_miura_cycl_stencil_03b( z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = (z_tracer_mflx_1_dsl - + z_tracer_mflx_2_dsl - + z_tracer_mflx_3_dsl)/float(3) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float( + 3 + ) return p_out_e @@ -38,8 +38,5 @@ def upwind_hflux_miura_cycl_stencil_03b( p_out_e: Field[[EdgeDim, KDim], float], ): _upwind_hflux_miura_cycl_stencil_03b( - z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl, - z_tracer_mflx_3_dsl, - out=(p_out_e) + z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, z_tracer_mflx_3_dsl, out=(p_out_e) ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py index 8ecbb1cb1a..aa972a9285 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.model.common.dimension import CellDim, KDim, EdgeDim, E2C +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator @@ -29,12 +29,22 @@ def _upwind_hflux_miura_stencil_01( cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) - + distv_bary_1 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) - + distv_bary_2 * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) - ) * p_mass_flx_e + p_out_e = ( + where( + cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0]) + ) + + distv_bary_1 + * where( + cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0]) + ) + + distv_bary_2 + * where( + cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0]) + ) + ) * p_mass_flx_e + + return p_out_e - return p_out_e @program def upwind_hflux_miura_stencil_01( @@ -55,5 +65,5 @@ def upwind_hflux_miura_stencil_01( distv_bary_2, p_mass_flx_e, cell_rel_idx_dsl, - out=(p_out_e) + out=(p_out_e), ) diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py index c9a59efb6e..4e4a3cabd7 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py @@ -15,10 +15,15 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import ( + btraj_dreg_stencil_01, +) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + _shape, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -31,9 +36,9 @@ def btraj_dreg_stencil_01_numpy( tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - lvn_sys_pos_true = np.where( tangent_orientation * p_vn >= 0.0, True, False) + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) @@ -61,7 +66,6 @@ def test_btraj_dreg_stencil_01(): tangent_orientation, lvn_sys_pos, offset_provider={}, - ) assert np.allclose(ref, lvn_sys_pos) diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py index 7e094d1fe2..d90c063df5 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py @@ -14,10 +14,17 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import ( + btraj_dreg_stencil_02, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -35,10 +42,10 @@ def btraj_dreg_stencil_02_numpy( e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) opt_famask_dsl = np.where( - traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), - int32(1), - int32(0) - ) + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), + int32(1), + int32(0), + ) return opt_famask_dsl @@ -49,14 +56,11 @@ def test_btraj_dreg_stencil_02(): p_vt = random_field(mesh, EdgeDim, KDim) edge_cell_length = np.asarray(mesh.e2c, dtype=float) edge_cell_length_new = as_1D_sparse_field(edge_cell_length, ECDim) - p_dt= 1.0 + p_dt = 1.0 opt_famask_dsl = zero_field(mesh, EdgeDim, KDim, dtype=int32) ref = btraj_dreg_stencil_02_numpy( - np.asarray(p_vn), - np.asarray(p_vt), - np.asarray(edge_cell_length), - p_dt + np.asarray(p_vn), np.asarray(p_vt), np.asarray(edge_cell_length), p_dt ) btraj_dreg_stencil_02( diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py index c55ae59395..42265f5486 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py @@ -16,10 +16,17 @@ from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 -from icon4py.model.common.dimension import KDim, E2CDim, ECDim, EdgeDim, CellDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field, constant_field +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import ( + btraj_dreg_stencil_03, +) +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + as_1D_sparse_field, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -38,8 +45,8 @@ def btraj_dreg_stencil_03_numpy( pos_on_tplane_e_2_y: np.array, primal_normal_cell_x: np.array, primal_normal_cell_y: np.array, - dual_normal_cell_x: np.array, - dual_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, lvn_sys_pos: np.array, p_dt: float, ) -> tuple[np.array]: @@ -163,7 +170,7 @@ def test_btraj_dreg_stencil_03(): primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) dual_normal_cell_y = random_field(mesh, EdgeDim, E2CDim) dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) - lvn_sys_pos = constant_field(mesh, True, EdgeDim, KDim, dtype=bool) + lvn_sys_pos = constant_field(mesh, True, EdgeDim, KDim, dtype=bool) p_dt = 2.0 p_cell_idx = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) @@ -189,7 +196,6 @@ def test_btraj_dreg_stencil_03(): p_coords_dreg_v_2_lat_dsl_ref, p_coords_dreg_v_3_lat_dsl_ref, p_coords_dreg_v_4_lat_dsl_ref, - ) = btraj_dreg_stencil_03_numpy( np.asarray(p_vn), np.asarray(p_vt), @@ -211,7 +217,6 @@ def test_btraj_dreg_stencil_03(): p_dt, ) - btraj_dreg_stencil_03( p_vn, p_vt, diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py index 5c91af2e45..9f11ad4f54 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py @@ -12,17 +12,22 @@ # 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.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( divide_flux_area_list_stencil_01, ) -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim - -from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + # FUNCTIONS # Checking turn when travelling along three points, used to check whether lines inters. def ccw( @@ -47,6 +52,7 @@ def ccw( ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise return ccw_out + # Checks whether two lines intersect def lintersect( line1_p1_lon: np.array, @@ -93,6 +99,7 @@ def lintersect( return lintersect_out + # Compute intersection point of two lines in 2D def line_intersect( line1_p1_lon: np.array, @@ -207,11 +214,15 @@ def divide_flux_area_list_stencil_01_numpy( ) lvn_sys_pos = np.where( - (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, True, False + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + True, + False, ) famask_bool = np.where(famask_int == int32(1), True, False) # ------------------------------------------------- Case 1 - mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + mask_case1 = np.logical_and.reduce( + [lintersect_line1, lintersect_line2, famask_bool] + ) ps1_x, ps1_y = line_intersect( fl_line_p1_lon, fl_line_p1_lat, @@ -234,76 +245,138 @@ def divide_flux_area_list_stencil_01_numpy( ) # Case 1 - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = np.where( + mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = np.where( + mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = np.where( - mask_case1, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = np.where( - mask_case1, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, ) dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + mask_case1, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) dreg_patch0_4_lat_dsl = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + mask_case1, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 1 - patch 1 dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0) - dreg_patch1_2_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0) - dreg_patch1_3_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0) - dreg_patch1_3_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) # Case 1 - patch 2 dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0) - dreg_patch2_2_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0) - dreg_patch2_3_lon_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0) - dreg_patch2_3_lat_vmask = np.where(mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) # ------------------------------------------------- Case 2a - mask_case2a = np.logical_and.reduce([lintersect_line1, np.logical_not(lintersect_line2), famask_bool]) + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] + ) # Case 2a - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = np.where( - mask_case2a, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = np.where( - mask_case2a, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl ) - dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = np.where( - mask_case2a, np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl + mask_case2a, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, ) dreg_patch0_4_lat_dsl = np.where( - mask_case2a, np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl + mask_case2a, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, ) # Case 2a - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_4_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask + ) dreg_patch1_2_lon_vmask = np.where( - mask_case2a, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask + mask_case2a, + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, ) dreg_patch1_2_lat_vmask = np.where( - mask_case2a, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask + mask_case2a, + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, ) dreg_patch1_3_lon_vmask = np.where( - mask_case2a, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, ) dreg_patch1_3_lat_vmask = np.where( - mask_case2a, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, ) # Case 2a - patch 2 dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) @@ -316,10 +389,16 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) # -------------------------------------------------- Case 2b - mask_case2b = np.logical_and.reduce([lintersect_line2, np.logical_not(lintersect_line1), famask_bool]) + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] + ) # Case 2b - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = np.where( + mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = np.where( + mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = np.where( mask_case2b, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -352,21 +431,37 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch1_4_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) # Case 2b - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_4_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask + ) dreg_patch2_2_lon_vmask = np.where( - mask_case2b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, ) dreg_patch2_2_lat_vmask = np.where( - mask_case2b, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, ) dreg_patch2_3_lon_vmask = np.where( - mask_case2b, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask + mask_case2b, + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, ) dreg_patch2_3_lat_vmask = np.where( - mask_case2b, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask + mask_case2b, + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, ) # flux area edge 1 and 2 @@ -403,8 +498,12 @@ def divide_flux_area_list_stencil_01_numpy( tri_line1_p2_lat, ) # Case 3a - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_1_lon_dsl = np.where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = np.where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) dreg_patch0_2_lon_dsl = np.where( mask_case3a, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -428,21 +527,37 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch0_4_lat_dsl, ) # Case 3a - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_1_lon_vmask = np.where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) dreg_patch1_2_lon_vmask = np.where( - mask_case3a, np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask + mask_case3a, + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, ) dreg_patch1_2_lat_vmask = np.where( - mask_case3a, np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask + mask_case3a, + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask ) - dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) dreg_patch1_4_lon_vmask = np.where( - mask_case3a, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, ) dreg_patch1_4_lat_vmask = np.where( - mask_case3a, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, ) # Case 3a - patch 2 dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) @@ -478,21 +593,37 @@ def divide_flux_area_list_stencil_01_numpy( tri_line2_p2_lat, ) # Case 3b - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_1_lon_dsl = np.where( + mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl + ) + dreg_patch0_1_lat_dsl = np.where( + mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl + ) + dreg_patch0_4_lon_dsl = np.where( + mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl + ) dreg_patch0_2_lon_dsl = np.where( - mask_case3b, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, ) dreg_patch0_2_lat_dsl = np.where( - mask_case3b, np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, ) dreg_patch0_3_lon_dsl = np.where( - mask_case3b, np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl + mask_case3b, + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, ) dreg_patch0_3_lat_dsl = np.where( - mask_case3b, np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl + mask_case3b, + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, ) # Case 3b - patch 1 dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) @@ -504,31 +635,49 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_1_lon_vmask = np.where( + mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) dreg_patch2_2_lon_vmask = np.where( - mask_case3b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, ) dreg_patch2_2_lat_vmask = np.where( - mask_case3b, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask ) - dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) dreg_patch2_4_lon_vmask = np.where( - mask_case3b, np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask + mask_case3b, + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, ) dreg_patch2_4_lat_vmask = np.where( - mask_case3b, np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask + mask_case3b, + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, ) # --------------------------------------------- Case 4 # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions # can't be overwritten by this new condition. - indices_previously_matched = ( - np.logical_or.reduce([mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1]) + indices_previously_matched = np.logical_or.reduce( + [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] ) # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = np.logical_and.reduce([famask_bool, np.logical_not(indices_previously_matched)]) + mask_case4 = np.logical_and.reduce( + [famask_bool, np.logical_not(indices_previously_matched)] + ) # Case 4 - patch 0 - no change # Case 4 - patch 1 dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) @@ -548,7 +697,6 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) - return ( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, @@ -576,6 +724,7 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask, ) + def test_divide_flux_area_list_stencil_01(): mesh = SimpleMesh() @@ -611,30 +760,32 @@ def test_divide_flux_area_list_stencil_01(): dreg_patch2_4_lon_vmask = zero_field(mesh, EdgeDim, KDim) dreg_patch2_4_lat_vmask = zero_field(mesh, EdgeDim, KDim) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ref_12, - ref_13, - ref_14, - ref_15, - ref_16, - ref_17, - ref_18, - ref_19, - ref_20, - ref_21, - ref_22, - ref_23, - ref_24) = divide_flux_area_list_stencil_01_numpy( + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20, + ref_21, + ref_22, + ref_23, + ref_24, + ) = divide_flux_area_list_stencil_01_numpy( mesh.e2c, np.asarray(famask_int), np.asarray(p_vn), diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py index 45916e8460..e920f1d738 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.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.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( divide_flux_area_list_stencil_02, ) -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, E2CDim, ECDim - -from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask, as_1D_sparse_field +from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -105,32 +109,31 @@ def divide_flux_area_list_stencil_02_numpy( dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat - # Store global index of the underlying grid cell # Adapt dimensions to fit ofr multiple levels butterfly_idx_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape + np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape ) butterfly_idx_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape + np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape ) butterfly_idx_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape + np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape ) butterfly_idx_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape + np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape ) butterfly_blk_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape + np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape ) butterfly_blk_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape + np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape ) butterfly_blk_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape + np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape ) butterfly_blk_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape + np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape ) patch1_cell_idx_vmask = np.where( famask_bool, @@ -189,7 +192,7 @@ def test_divide_flux_area_list_stencil_02(): bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) bf_cc_patch2_lat = random_field(mesh, EdgeDim, E2CDim) bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) - butterfly_idx_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) butterfly_idx_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) butterfly_blk_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) butterfly_blk_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) @@ -218,26 +221,28 @@ def test_divide_flux_area_list_stencil_02(): patch2_cell_idx_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) patch2_cell_blk_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ref_12, - ref_13, - ref_14, - ref_15, - ref_16, - ref_17, - ref_18, - ref_19, - ref_20) = divide_flux_area_list_stencil_02_numpy( + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20, + ) = divide_flux_area_list_stencil_02_numpy( mesh.e2c, np.asarray(famask_int), np.asarray(p_vn), @@ -327,7 +332,7 @@ def test_divide_flux_area_list_stencil_02(): assert np.allclose(dreg_patch2_3_lat_vmask, ref_14) assert np.allclose(dreg_patch2_4_lon_vmask, ref_15) assert np.allclose(dreg_patch2_4_lat_vmask, ref_16) - assert np.allclose(patch1_cell_idx_vmask, ref_17) - assert np.allclose(patch1_cell_blk_vmask, ref_18) - assert np.allclose(patch2_cell_idx_vmask, ref_19) - assert np.allclose(patch2_cell_blk_vmask, ref_20) + assert np.allclose(patch1_cell_idx_vmask, ref_17) + assert np.allclose(patch1_cell_blk_vmask, ref_18) + assert np.allclose(patch2_cell_idx_vmask, ref_19) + assert np.allclose(patch2_cell_blk_vmask, ref_20) diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py index be97aa4b4f..69412e0b7d 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py @@ -15,10 +15,15 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import ( + face_val_ppm_stencil_01, +) from icon4py.model.common.dimension import CellDim, KDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + _shape, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py index 5eed1821d8..461c0bffd0 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py @@ -15,9 +15,10 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import ( + face_val_ppm_stencil_02, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import _shape, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py index fa827a12b3..a828e5897f 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import ( + face_val_ppm_stencil_02a, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py index 7a87257d42..6b1b867879 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import ( + face_val_ppm_stencil_02b, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py index f4787f290a..a02c6d25d6 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import ( + face_val_ppm_stencil_02c, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py index d83b6389c0..7dd8ca373f 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import ( + face_val_ppm_stencil_05, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py index 846258a19f..024dd0ae90 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -15,10 +15,16 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import hflux_ffsl_hybrid_stencil_01a -from icon4py.model.common.dimension import CellDim, KDim, EdgeDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( + hflux_ffsl_hybrid_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -47,29 +53,79 @@ def hflux_ffsl_hybrid_stencil_01a_numpy( patch0_cell_rel_idx_dsl: np.ndarray, ): - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] - z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] - z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] - z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] - z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] - z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] p_out_e_hybrid_1a = ( - np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) * z_quad_vector_sum0_1 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) * z_quad_vector_sum0_2 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) * z_quad_vector_sum0_3 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_4_e2c[:, 1], z_lsq_coeff_4_e2c[:, 0]) * z_quad_vector_sum0_4 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_5_e2c[:, 1], z_lsq_coeff_5_e2c[:, 0]) * z_quad_vector_sum0_5 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_6_e2c[:, 1], z_lsq_coeff_6_e2c[:, 0]) * z_quad_vector_sum0_6 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_7_e2c[:, 1], z_lsq_coeff_7_e2c[:, 0]) * z_quad_vector_sum0_7 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_8_e2c[:, 1], z_lsq_coeff_8_e2c[:, 0]) * z_quad_vector_sum0_8 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_9_e2c[:, 1], z_lsq_coeff_9_e2c[:, 0]) * z_quad_vector_sum0_9 - + np.where(patch0_cell_rel_idx_dsl == int32(1), z_lsq_coeff_10_e2c[:, 1], z_lsq_coeff_10_e2c[:, 0]) * z_quad_vector_sum0_10 - ) + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum0_1 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum0_2 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum0_3 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum0_4 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum0_5 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum0_6 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum0_7 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum0_8 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum0_9 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum0_10 + ) return p_out_e_hybrid_1a @@ -150,7 +206,6 @@ def test_hflux_ffsl_hybrid_stencil_01a(): offset_provider={ "E2C": mesh.get_e2c_offset_provider(), }, - ) assert np.allclose(p_out_e_hybrid_1a, ref) diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py index dcde5d491a..ecfe1750a5 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -15,10 +15,16 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import hflux_ffsl_hybrid_stencil_02 -from icon4py.model.common.dimension import KDim, EdgeDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( + hflux_ffsl_hybrid_stencil_02, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py index 9da194cc6f..85f9ca14c7 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py @@ -15,10 +15,16 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import hflx_limiter_mo_stencil_01a -from icon4py.model.common.dimension import KDim, EdgeDim, CellDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( + hflx_limiter_mo_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -48,7 +54,6 @@ def test_hflx_limiter_mo_stencil_01a(): z_mflx_low = zero_field(mesh, EdgeDim, KDim) z_anti = zero_field(mesh, EdgeDim, KDim) - ref_1, ref_2 = hflx_limiter_mo_stencil_01a_numpy( mesh.e2c, np.asarray(p_mflx_tracer_h), diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py index f4088296f3..8f7f4f11af 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py @@ -16,10 +16,17 @@ from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import hflx_limiter_mo_stencil_01b -from icon4py.model.common.dimension import C2EDim, CEDim, KDim, EdgeDim, CellDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, constant_field, as_1D_sparse_field +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( + hflx_limiter_mo_stencil_01b, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + as_1D_sparse_field, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -38,20 +45,21 @@ def hflx_limiter_mo_stencil_01b_numpy( zero_array = np.zeros(p_rhodz_now.shape) - z_mflx_anti_1 = ( p_dtime * geofac_div[:, 0] / p_rhodz_new - * z_anti_c2e[:, 0] ) - z_mflx_anti_2 = ( p_dtime * geofac_div[:, 1] / p_rhodz_new - * z_anti_c2e[:, 1] ) - z_mflx_anti_3 = ( p_dtime * geofac_div[:, 2] / p_rhodz_new - * z_anti_c2e[:, 2] ) + z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] + z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] + z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] - z_mflx_anti_in = -1.0 * ( np.minimum(zero_array, z_mflx_anti_1) - + np.minimum(zero_array, z_mflx_anti_2) - + np.minimum(zero_array, z_mflx_anti_3) ) + z_mflx_anti_in = -1.0 * ( + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) + ) - z_mflx_anti_out = ( np.maximum(zero_array, z_mflx_anti_1) - + np.maximum(zero_array, z_mflx_anti_2) - + np.maximum(zero_array, z_mflx_anti_3) ) + z_mflx_anti_out = ( + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) + ) z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) @@ -67,6 +75,7 @@ def hflx_limiter_mo_stencil_01b_numpy( z_tracer_min, ) + def test_hflx_limiter_mo_stencil_01b(): mesh = SimpleMesh() @@ -92,7 +101,7 @@ def test_hflx_limiter_mo_stencil_01b(): np.asarray(z_mflx_low), np.asarray(z_anti), np.asarray(p_cc), - np.asarray(p_dtime) + np.asarray(p_dtime), ) hflx_limiter_mo_stencil_01b( @@ -114,8 +123,8 @@ def test_hflx_limiter_mo_stencil_01b(): }, ) - assert np.allclose(z_mflx_anti_in, ref_1) - assert np.allclose(z_mflx_anti_out, ref_2) + assert np.allclose(z_mflx_anti_in, ref_1) + assert np.allclose(z_mflx_anti_out, ref_2) assert np.allclose(z_tracer_new_low, ref_3) - assert np.allclose(z_tracer_max, ref_4) - assert np.allclose(z_tracer_min, ref_5) + assert np.allclose(z_tracer_max, ref_4) + assert np.allclose(z_tracer_min, ref_5) diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py index 66cbb5212f..55a025afb5 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -18,8 +18,11 @@ hflx_limiter_mo_stencil_02, ) from icon4py.model.common.dimension import CellDim, KDim - -from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py index 5264b1e4b4..9677974bc8 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -17,7 +17,6 @@ hflx_limiter_mo_stencil_03_min_max, ) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py index bd04b9b914..d2d29d7ca6 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -17,7 +17,6 @@ hflx_limiter_mo_stencil_04, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py index 19404c50e0..3533162d39 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -18,8 +18,11 @@ hflx_limiter_pd_stencil_01, ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py index dcef77e331..b127080555 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -17,7 +17,6 @@ hflx_limiter_pd_stencil_02, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - from icon4py.model.common.test_utils.helpers import constant_field, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py index 4a04d00962..458119d1b1 100644 --- a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py @@ -14,10 +14,14 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 +from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import ( + hor_adv_stencil_01, +) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py index 12c75d4184..95ced3a802 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -12,15 +12,19 @@ # 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.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( prep_gauss_quadrature_c_list_stencil, ) -from icon4py.model.common.dimension import EdgeDim, CellDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field, constant_field +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -99,7 +103,7 @@ def prep_gauss_quadrature_c_list_stencil_numpy( p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) - wgt_t_detjac_1 =np.where( + wgt_t_detjac_1 = np.where( famask_bool, dbl_eps + z_wgt_1 @@ -123,7 +127,7 @@ def prep_gauss_quadrature_c_list_stencil_numpy( ), 0.0, ) - wgt_t_detjac_2 =np.where( + wgt_t_detjac_2 = np.where( famask_bool, dbl_eps + z_wgt_2 @@ -147,7 +151,7 @@ def prep_gauss_quadrature_c_list_stencil_numpy( ), 0.0, ) - wgt_t_detjac_3 =np.where( + wgt_t_detjac_3 = np.where( famask_bool, dbl_eps + z_wgt_3 @@ -170,8 +174,8 @@ def prep_gauss_quadrature_c_list_stencil_numpy( ) ), 0.0, - ) - wgt_t_detjac_4 =np.where( + ) + wgt_t_detjac_4 = np.where( famask_bool, dbl_eps + z_wgt_4 @@ -318,11 +322,12 @@ def prep_gauss_quadrature_c_list_stencil_numpy( p_dreg_area, ) + def test_prep_gauss_quadrature_c_list_stencil(): mesh = SimpleMesh() - famask_int = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) - p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + famask_int = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) @@ -330,7 +335,7 @@ def test_prep_gauss_quadrature_c_list_stencil(): p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) - shape_func_1_1 = 0.001 + shape_func_1_1 = 0.001 shape_func_2_1 = 0.001 shape_func_3_1 = 0.001 shape_func_4_1 = 0.001 @@ -346,11 +351,11 @@ def test_prep_gauss_quadrature_c_list_stencil(): shape_func_2_4 = 0.001 shape_func_3_4 = 0.001 shape_func_4_4 = 0.001 - zeta_1 = 0.002 + zeta_1 = 0.002 zeta_2 = 0.002 zeta_3 = 0.002 zeta_4 = 0.002 - eta_1 = 0.5 + eta_1 = 0.5 eta_2 = 0.5 eta_3 = 0.5 eta_4 = 0.5 @@ -361,7 +366,7 @@ def test_prep_gauss_quadrature_c_list_stencil(): dbl_eps = np.float64(0.1) eps = 0.1 p_dreg_area_in = random_field(mesh, EdgeDim, KDim) - p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) @@ -373,57 +378,59 @@ def test_prep_gauss_quadrature_c_list_stencil(): p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) p_dreg_area = zero_field(mesh, EdgeDim, KDim) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11) = prep_gauss_quadrature_c_list_stencil_numpy( - np.asarray(famask_int), - np.asarray(p_coords_dreg_v_1_x), - np.asarray(p_coords_dreg_v_2_x), - np.asarray(p_coords_dreg_v_3_x), - np.asarray(p_coords_dreg_v_4_x), - np.asarray(p_coords_dreg_v_1_y), - np.asarray(p_coords_dreg_v_2_y), - np.asarray(p_coords_dreg_v_3_y), - np.asarray(p_coords_dreg_v_4_y), - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, - np.asarray(p_dreg_area_in), + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ) = prep_gauss_quadrature_c_list_stencil_numpy( + np.asarray(famask_int), + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + np.asarray(p_dreg_area_in), ) prep_gauss_quadrature_c_list_stencil( diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py index 032041b648..77e574074a 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py @@ -17,9 +17,12 @@ from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( prep_gauss_quadrature_c_stencil, ) -from icon4py.model.common.dimension import EdgeDim, CellDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -268,7 +271,9 @@ def prep_gauss_quadrature_c_stencil_numpy( z_area = p_quad_vector_sum_1 p_dreg_area_out = np.where( - z_area >= 0.0, np.maximum(eps, np.absolute(z_area)), -np.maximum(eps, np.absolute(z_area)) + z_area >= 0.0, + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), ) return ( p_quad_vector_sum_1, @@ -284,10 +289,11 @@ def prep_gauss_quadrature_c_stencil_numpy( p_dreg_area_out, ) + def test_prep_gauss_quadrature_c_stencil(): mesh = SimpleMesh() - p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) @@ -295,7 +301,7 @@ def test_prep_gauss_quadrature_c_stencil(): p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) - shape_func_1_1 = 0.001 + shape_func_1_1 = 0.001 shape_func_2_1 = 0.001 shape_func_3_1 = 0.001 shape_func_4_1 = 0.001 @@ -311,11 +317,11 @@ def test_prep_gauss_quadrature_c_stencil(): shape_func_2_4 = 0.001 shape_func_3_4 = 0.001 shape_func_4_4 = 0.001 - zeta_1 = 0.002 + zeta_1 = 0.002 zeta_2 = 0.002 zeta_3 = 0.002 zeta_4 = 0.002 - eta_1 = 0.5 + eta_1 = 0.5 eta_2 = 0.5 eta_3 = 0.5 eta_4 = 0.5 @@ -325,7 +331,7 @@ def test_prep_gauss_quadrature_c_stencil(): wgt_eta_2 = 0.007 dbl_eps = np.float64(0.1) eps = 0.1 - p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) @@ -337,55 +343,57 @@ def test_prep_gauss_quadrature_c_stencil(): p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) p_dreg_area_out = zero_field(mesh, EdgeDim, KDim) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11) = prep_gauss_quadrature_c_stencil_numpy( - np.asarray(p_coords_dreg_v_1_x), - np.asarray(p_coords_dreg_v_2_x), - np.asarray(p_coords_dreg_v_3_x), - np.asarray(p_coords_dreg_v_4_x), - np.asarray(p_coords_dreg_v_1_y), - np.asarray(p_coords_dreg_v_2_y), - np.asarray(p_coords_dreg_v_3_y), - np.asarray(p_coords_dreg_v_4_y), - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ) = prep_gauss_quadrature_c_stencil_numpy( + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, ) prep_gauss_quadrature_c_stencil( diff --git a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py index 3ac9f1bd8d..2bf78bc992 100644 --- a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 +from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import ( + rbf_intp_edge_stencil_01, +) from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py index af750be274..916e5df693 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py @@ -18,14 +18,17 @@ recon_lsq_cell_c_stencil, ) from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def recon_lsq_cell_c_stencil_numpy( c2e2c2e2c: np.ndarray, - p_cc: np.ndarray, + p_cc: np.ndarray, lsq_qtmat_c_1: np.ndarray, lsq_qtmat_c_2: np.ndarray, lsq_qtmat_c_3: np.ndarray, @@ -88,10 +91,10 @@ def recon_lsq_cell_c_stencil_numpy( lsq_moments_6: np.ndarray, lsq_moments_7: np.ndarray, lsq_moments_8: np.ndarray, - lsq_moments_9: np.ndarray + lsq_moments_9: np.ndarray, ) -> tuple[np.ndarray]: p_cc_e = np.expand_dims(p_cc, axis=-1) -# n_diff = p_cc[c2e2c] - p_cc_e + # n_diff = p_cc[c2e2c] - p_cc_e lsq_rmat_rdiag_c_1 = np.expand_dims(lsq_rmat_rdiag_c_1, axis=-1) lsq_rmat_rdiag_c_2 = np.expand_dims(lsq_rmat_rdiag_c_2, axis=-1) @@ -102,15 +105,15 @@ def recon_lsq_cell_c_stencil_numpy( lsq_rmat_rdiag_c_7 = np.expand_dims(lsq_rmat_rdiag_c_7, axis=-1) lsq_rmat_rdiag_c_8 = np.expand_dims(lsq_rmat_rdiag_c_8, axis=-1) lsq_rmat_rdiag_c_9 = np.expand_dims(lsq_rmat_rdiag_c_9, axis=-1) -# lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) -# lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) -# lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) -# lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) -# lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) -# lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) -# lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) -# lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) -# lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) + # lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) + # lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) + # lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) + # lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) + # lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) + # lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) + # lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) + # lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) + # lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) @@ -201,33 +204,33 @@ def recon_lsq_cell_c_stencil_numpy( lsq_rmat_utri_c_34 = np.broadcast_to(lsq_rmat_utri_c_34, p_cc.shape) lsq_rmat_utri_c_35 = np.broadcast_to(lsq_rmat_utri_c_35, p_cc.shape) lsq_rmat_utri_c_36 = np.broadcast_to(lsq_rmat_utri_c_36, p_cc.shape) -# lsq_rmat_utri_c_9 lsq_qtmat_c_1 = nplsq_rmat_utri_c_9.broadcast_to(lsq_qtmat_c_1, (CECECDim, KDim)) -# lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, (CECECDim, KDim)) -# lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, (CECECDim, KDim)) -# lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, (CECECDim, KDim)) -# lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, (CECECDim, KDim)) -# lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, (CECECDim, KDim)) -# lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, (CECECDim, KDim)) -# lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, (CECECDim, KDim)) -# lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, (CECECDim, KDim)) -# lsq_qtmat_c_1 = np.broadcast_to(lsq_qtmat_c_1, p_cc.shape) -# lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, p_cc.shape) -# lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, p_cc.shape) -# lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, p_cc.shape) -# lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, p_cc.shape) -# lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, p_cc.shape) -# lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, p_cc.shape) -# lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, p_cc.shape) -# lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, p_cc.shape) -# lsq_qtmat_c_1 = np.repeat(lsq_qtmat_c_1[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_2 = np.repeat(lsq_qtmat_c_2[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_3 = np.repeat(lsq_qtmat_c_3[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_4 = np.repeat(lsq_qtmat_c_4[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_5 = np.repeat(lsq_qtmat_c_5[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_6 = np.repeat(lsq_qtmat_c_6[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_7 = np.repeat(lsq_qtmat_c_7[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_8 = np.repeat(lsq_qtmat_c_8[:, np.newaxis], p_cc.shape[1], axis=1) -# lsq_qtmat_c_9 = np.repeat(lsq_qtmat_c_9[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_rmat_utri_c_9 lsq_qtmat_c_1 = nplsq_rmat_utri_c_9.broadcast_to(lsq_qtmat_c_1, (CECECDim, KDim)) + # lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, (CECECDim, KDim)) + # lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, (CECECDim, KDim)) + # lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, (CECECDim, KDim)) + # lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, (CECECDim, KDim)) + # lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, (CECECDim, KDim)) + # lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, (CECECDim, KDim)) + # lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, (CECECDim, KDim)) + # lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, (CECECDim, KDim)) + # lsq_qtmat_c_1 = np.broadcast_to(lsq_qtmat_c_1, p_cc.shape) + # lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, p_cc.shape) + # lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, p_cc.shape) + # lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, p_cc.shape) + # lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, p_cc.shape) + # lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, p_cc.shape) + # lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, p_cc.shape) + # lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, p_cc.shape) + # lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, p_cc.shape) + # lsq_qtmat_c_1 = np.repeat(lsq_qtmat_c_1[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_2 = np.repeat(lsq_qtmat_c_2[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_3 = np.repeat(lsq_qtmat_c_3[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_4 = np.repeat(lsq_qtmat_c_4[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_5 = np.repeat(lsq_qtmat_c_5[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_6 = np.repeat(lsq_qtmat_c_6[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_7 = np.repeat(lsq_qtmat_c_7[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_8 = np.repeat(lsq_qtmat_c_8[:, np.newaxis], p_cc.shape[1], axis=1) + # lsq_qtmat_c_9 = np.repeat(lsq_qtmat_c_9[:, np.newaxis], p_cc.shape[1], axis=1) lsq_qtmat_c_9 = np.expand_dims(lsq_qtmat_c_9, axis=-1) lsq_qtmat_c_8 = np.expand_dims(lsq_qtmat_c_8, axis=-1) lsq_qtmat_c_7 = np.expand_dims(lsq_qtmat_c_7, axis=-1) @@ -239,20 +242,19 @@ def recon_lsq_cell_c_stencil_numpy( lsq_qtmat_c_1 = np.expand_dims(lsq_qtmat_c_1, axis=-1) p_coeff_10 = lsq_rmat_rdiag_c_9 * ( - lsq_qtmat_c_9[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_9[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_9[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_9[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_9[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_9[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_9[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_9[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_9[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_9[:, 0] * (p_cc_e[:, 0] - p_cc) + + lsq_qtmat_c_9[:, 1] * (p_cc_e[:, 1] - p_cc) + + lsq_qtmat_c_9[:, 2] * (p_cc_e[:, 2] - p_cc) + + lsq_qtmat_c_9[:, 3] * (p_cc_e[:, 3] - p_cc) + + lsq_qtmat_c_9[:, 4] * (p_cc_e[:, 4] - p_cc) + + lsq_qtmat_c_9[:, 5] * (p_cc_e[:, 5] - p_cc) + + lsq_qtmat_c_9[:, 6] * (p_cc_e[:, 6] - p_cc) + + lsq_qtmat_c_9[:, 7] * (p_cc_e[:, 7] - p_cc) + + lsq_qtmat_c_9[:, 8] * (p_cc_e[:, 8] - p_cc) ) - p_coeff_9 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_8[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_8[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_8[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_8[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_8[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -261,11 +263,11 @@ def recon_lsq_cell_c_stencil_numpy( + lsq_qtmat_c_8[:, 6] * (p_cc_e[:, 6] - p_cc) + lsq_qtmat_c_8[:, 7] * (p_cc_e[:, 7] - p_cc) + lsq_qtmat_c_8[:, 8] * (p_cc_e[:, 8] - p_cc) - - lsq_rmat_utri_c_1 * p_coeff_10 - ) + - lsq_rmat_utri_c_1 * p_coeff_10 + ) p_coeff_8 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_7[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_7[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_7[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_7[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_7[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -274,11 +276,11 @@ def recon_lsq_cell_c_stencil_numpy( + lsq_qtmat_c_7[:, 6] * (p_cc_e[:, 6] - p_cc) + lsq_qtmat_c_7[:, 7] * (p_cc_e[:, 7] - p_cc) + lsq_qtmat_c_7[:, 8] * (p_cc_e[:, 8] - p_cc) - - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) - ) + - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) + ) p_coeff_7 = lsq_rmat_rdiag_c_6 * ( - lsq_qtmat_c_6[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_6[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_6[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_6[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_6[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -295,7 +297,7 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_6 = lsq_rmat_rdiag_c_5 * ( - lsq_qtmat_c_5[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_5[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_5[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_5[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_5[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -313,7 +315,7 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_5 = lsq_rmat_rdiag_c_4 * ( - lsq_qtmat_c_4[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_4[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_4[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_4[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_4[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -322,7 +324,7 @@ def recon_lsq_cell_c_stencil_numpy( + lsq_qtmat_c_4[:, 6] * (p_cc_e[:, 6] - p_cc) + lsq_qtmat_c_4[:, 7] * (p_cc_e[:, 7] - p_cc) + lsq_qtmat_c_4[:, 8] * (p_cc_e[:, 8] - p_cc) - - ( + - ( lsq_rmat_utri_c_11 * p_coeff_6 + lsq_rmat_utri_c_12 * p_coeff_7 + lsq_rmat_utri_c_13 * p_coeff_8 @@ -332,7 +334,7 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_4 = lsq_rmat_rdiag_c_3 * ( - lsq_qtmat_c_3[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_3[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_3[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_3[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_3[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -352,7 +354,7 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_3 = lsq_rmat_rdiag_c_2 * ( - lsq_qtmat_c_2[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_2[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_2[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_2[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_2[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -373,7 +375,7 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_2 = lsq_rmat_rdiag_c_1 * ( - lsq_qtmat_c_1[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_1[:, 0] * (p_cc_e[:, 0] - p_cc) + lsq_qtmat_c_1[:, 1] * (p_cc_e[:, 1] - p_cc) + lsq_qtmat_c_1[:, 2] * (p_cc_e[:, 2] - p_cc) + lsq_qtmat_c_1[:, 3] * (p_cc_e[:, 3] - p_cc) @@ -418,6 +420,7 @@ def recon_lsq_cell_c_stencil_numpy( p_coeff_10, ) + def test_recon_lsq_cell_c_stencil(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) @@ -440,7 +443,7 @@ def test_recon_lsq_cell_c_stencil(): lsq_qtmat_c_8_field = as_1D_sparse_field(lsq_qtmat_c_8, CECECDim) lsq_qtmat_c_9_field = as_1D_sparse_field(lsq_qtmat_c_9, CECECDim) lsq_rmat_rdiag_c_1 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_2 = random_field(mesh, CellDim) + lsq_rmat_rdiag_c_2 = random_field(mesh, CellDim) lsq_rmat_rdiag_c_3 = random_field(mesh, CellDim) lsq_rmat_rdiag_c_4 = random_field(mesh, CellDim) lsq_rmat_rdiag_c_5 = random_field(mesh, CellDim) @@ -504,16 +507,18 @@ def test_recon_lsq_cell_c_stencil(): p_coeff_9 = zero_field(mesh, CellDim, KDim) p_coeff_10 = zero_field(mesh, CellDim, KDim) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10) = recon_lsq_cell_c_stencil_numpy( + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ) = recon_lsq_cell_c_stencil_numpy( mesh.c2e2c2e2c, np.asarray(p_cc), np.asarray(lsq_qtmat_c_1), @@ -578,7 +583,7 @@ def test_recon_lsq_cell_c_stencil(): np.asarray(lsq_moments_6), np.asarray(lsq_moments_7), np.asarray(lsq_moments_8), - np.asarray(lsq_moments_9) + np.asarray(lsq_moments_9), ) recon_lsq_cell_c_stencil( @@ -646,7 +651,7 @@ def test_recon_lsq_cell_c_stencil(): lsq_moments_7, lsq_moments_8, lsq_moments_9, - p_coeff_1, + p_coeff_1, p_coeff_2, p_coeff_3, p_coeff_4, @@ -658,7 +663,9 @@ def test_recon_lsq_cell_c_stencil(): p_coeff_10, offset_provider={ "C2E2C2E2C": mesh.get_c2e2c_offset_provider(), - "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), + "C2CECEC": StridedNeighborOffsetProvider( + CellDim, CECECDim, mesh.n_c2e2c2e2c + ), }, ) co1 = np.asarray(p_coeff_1) diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py index 95d0116da1..8945975f34 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py @@ -18,13 +18,17 @@ recon_lsq_cell_c_svd_stencil, ) from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + def recon_lsq_cell_c_svd_stencil_numpy( c2e2c2e2c: np.ndarray, - p_cc: np.ndarray, + p_cc: np.ndarray, lsq_pseudoinv_1: np.ndarray, lsq_pseudoinv_2: np.ndarray, lsq_pseudoinv_3: np.ndarray, @@ -42,7 +46,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( lsq_moments_6: np.ndarray, lsq_moments_7: np.ndarray, lsq_moments_8: np.ndarray, - lsq_moments_9: np.ndarray + lsq_moments_9: np.ndarray, ) -> tuple[np.ndarray]: p_cc_e = np.expand_dims(p_cc, axis=-1) @@ -75,20 +79,19 @@ def recon_lsq_cell_c_svd_stencil_numpy( lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) p_coeff_10_dsl = ( - lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) ) - p_coeff_9_dsl = ( - lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -97,10 +100,10 @@ def recon_lsq_cell_c_svd_stencil_numpy( + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + ) p_coeff_8_dsl = ( - lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -109,10 +112,10 @@ def recon_lsq_cell_c_svd_stencil_numpy( + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + ) p_coeff_7_dsl = ( - lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -124,7 +127,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) p_coeff_6_dsl = ( - lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -136,7 +139,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) p_coeff_5_dsl = ( - lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -148,7 +151,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) p_coeff_4_dsl = ( - lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -160,7 +163,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) p_coeff_3_dsl = ( - lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -172,7 +175,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) p_coeff_2_dsl = ( - lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) @@ -207,6 +210,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( p_coeff_10_dsl, ) + def test_recon_lsq_cell_c_svd_stencil(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) @@ -248,16 +252,18 @@ def test_recon_lsq_cell_c_svd_stencil(): p_coeff_9_dsl = zero_field(mesh, CellDim, KDim) p_coeff_10_dsl = zero_field(mesh, CellDim, KDim) - (ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10) = recon_lsq_cell_c_svd_stencil_numpy( + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ) = recon_lsq_cell_c_svd_stencil_numpy( mesh.c2e2c2e2c, np.asarray(p_cc), np.asarray(lsq_pseudoinv_1), @@ -277,7 +283,7 @@ def test_recon_lsq_cell_c_svd_stencil(): np.asarray(lsq_moments_6), np.asarray(lsq_moments_7), np.asarray(lsq_moments_8), - np.asarray(lsq_moments_9) + np.asarray(lsq_moments_9), ) recon_lsq_cell_c_svd_stencil( @@ -300,7 +306,7 @@ def test_recon_lsq_cell_c_svd_stencil(): lsq_moments_7, lsq_moments_8, lsq_moments_9, - p_coeff_1_dsl, + p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl, p_coeff_4_dsl, @@ -312,10 +318,12 @@ def test_recon_lsq_cell_c_svd_stencil(): p_coeff_10_dsl, offset_provider={ "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), - "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), + "C2CECEC": StridedNeighborOffsetProvider( + CellDim, CECECDim, mesh.n_c2e2c2e2c + ), }, ) -# co1 = np.asarray(p_coeff_1_dsl) + # co1 = np.asarray(p_coeff_1_dsl) co2 = np.asarray(p_coeff_2_dsl) co3 = np.asarray(p_coeff_3_dsl) co4 = np.asarray(p_coeff_4_dsl) @@ -325,7 +333,7 @@ def test_recon_lsq_cell_c_svd_stencil(): co8 = np.asarray(p_coeff_8_dsl) co9 = np.asarray(p_coeff_9_dsl) co10 = np.asarray(p_coeff_10_dsl) -# assert np.allclose(ref_1, co1) + # assert np.allclose(ref_1, co1) assert np.allclose(ref_10, co10) assert np.allclose(ref_2, co2) assert np.allclose(ref_3, co3) diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py index e31a49d40d..3f7dca6ede 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py @@ -18,8 +18,11 @@ recon_lsq_cell_l_svd_stencil, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim - -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_set_zero_c.py b/model/atmosphere/advection/tests/test_set_zero_c.py index 059e47adcd..f023b44e61 100644 --- a/model/atmosphere/advection/tests/test_set_zero_c.py +++ b/model/atmosphere/advection/tests/test_set_zero_c.py @@ -15,7 +15,6 @@ from icon4py.model.atmosphere.advection.set_zero_c import set_zero_c from icon4py.model.common.dimension import CellDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_set_zero_c_k.py b/model/atmosphere/advection/tests/test_set_zero_c_k.py index 6f26b960fd..e0f4afc4b4 100644 --- a/model/atmosphere/advection/tests/test_set_zero_c_k.py +++ b/model/atmosphere/advection/tests/test_set_zero_c_k.py @@ -15,7 +15,6 @@ from icon4py.model.atmosphere.advection.set_zero_c_k import set_zero_c_k from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py index 803c9e62eb..ce87c96606 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.model.atmosphere.advection.step_advection_stencil_01 import ( + step_advection_stencil_01, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py index bc60330323..fa15a23593 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.model.atmosphere.advection.step_advection_stencil_02 import ( + step_advection_stencil_02, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py index e69a82f6d0..02a75ffef1 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 +from icon4py.model.atmosphere.advection.step_advection_stencil_03 import ( + step_advection_stencil_03, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py index 36addc80e9..7a57b2adc4 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 +from icon4py.model.atmosphere.advection.step_advection_stencil_04 import ( + step_advection_stencil_04, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py index 3bb27e469c..6fff27fbce 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py @@ -12,14 +12,17 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( upwind_hflux_miura3_stencil_01, ) -from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import random_field, zero_field, random_mask +from icon4py.model.common.test_utils.helpers import ( + random_field, + random_mask, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -62,17 +65,71 @@ def upwind_hflux_miura3_stencil_01_numpy( z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] p_out_e_miura3 = ( - np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) * z_quad_vector_sum_1 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) * z_quad_vector_sum_2 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) * z_quad_vector_sum_3 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_4_e2c[:, 1], z_lsq_coeff_4_e2c[:, 0]) * z_quad_vector_sum_4 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_5_e2c[:, 1], z_lsq_coeff_5_e2c[:, 0]) * z_quad_vector_sum_5 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_6_e2c[:, 1], z_lsq_coeff_6_e2c[:, 0]) * z_quad_vector_sum_6 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_7_e2c[:, 1], z_lsq_coeff_7_e2c[:, 0]) * z_quad_vector_sum_7 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_8_e2c[:, 1], z_lsq_coeff_8_e2c[:, 0]) * z_quad_vector_sum_8 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_9_e2c[:, 1], z_lsq_coeff_9_e2c[:, 0]) * z_quad_vector_sum_9 - + np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_10_e2c[:, 1], z_lsq_coeff_10_e2c[:, 0]) * z_quad_vector_sum_10 - ) / z_dreg_area * p_mass_flx_e + ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum_1 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum_2 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum_3 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum_4 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum_5 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum_6 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum_7 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum_8 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum_9 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum_10 + ) + / z_dreg_area + * p_mass_flx_e + ) return p_out_e_miura3 @@ -102,7 +159,7 @@ def test_upwind_hflux_miura3_stencil_01(): z_quad_vector_sum_10 = random_field(mesh, EdgeDim, KDim) p_mass_flx_e = random_field(mesh, EdgeDim, KDim) z_dreg_area = random_field(mesh, EdgeDim, KDim) -# cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + # cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) p_out_e_miura3 = zero_field(mesh, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py index 1939b09564..ef9c229747 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -14,10 +14,16 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import upwind_hflux_miura_cycl_stencil_01 +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import ( + upwind_hflux_miura_cycl_stencil_01, +) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, random_mask +from icon4py.model.common.test_utils.helpers import ( + _shape, + random_field, + random_mask, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -35,10 +41,25 @@ def upwind_hflux_miura_cycl_stencil_01_numpy( z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] - z_tracer_mflx_dsl = ( np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_dsl_e2c[:, 1], z_lsq_coeff_1_dsl_e2c[:, 0]) - + distv_bary_1 * np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_dsl_e2c[:, 1], z_lsq_coeff_2_dsl_e2c[:, 0]) - + distv_bary_2 * np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_dsl_e2c[:, 1], z_lsq_coeff_3_dsl_e2c[:, 0]) - ) * p_mass_flx_e + z_tracer_mflx_dsl = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl_e2c[:, 1], + z_lsq_coeff_1_dsl_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl_e2c[:, 1], + z_lsq_coeff_2_dsl_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl_e2c[:, 1], + z_lsq_coeff_3_dsl_e2c[:, 0], + ) + ) * p_mass_flx_e return z_tracer_mflx_dsl diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py index a89ae88c34..b6d898e080 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -15,10 +15,16 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import upwind_hflux_miura_cycl_stencil_02 -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, CEDim, C2EDim - -from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field, as_1D_sparse_field +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import ( + upwind_hflux_miura_cycl_stencil_02, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + _shape, + as_1D_sparse_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -37,19 +43,22 @@ def upwind_hflux_miura_cycl_stencil_02_numpy( geofac_div = np.expand_dims(geofac_div, axis=-1) z_tracer_mflx_c2e = z_tracer_mflx[c2e] - z_rhofluxdiv_c_out = np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c + z_rhofluxdiv_c_out = ( + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) + if nsub == int32(1) + else z_rhofluxdiv_c + ) z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) - z_rho_new_dsl = ( z_rho_now - - z_dtsub * z_rhofluxdiv_c_out ) + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out - z_tracer_new_dsl = (( z_tracer_now - * z_rho_now - - z_dtsub * z_fluxdiv_c_dsl ) - / z_rho_new_dsl) + z_tracer_new_dsl = ( + z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl + ) / z_rho_new_dsl return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + def test_upwind_hflux_miura_cycl_stencil_02(): mesh = SimpleMesh() nsub = int32(1) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py index 52110a21a5..6e55efa77e 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import upwind_hflux_miura_cycl_stencil_03a +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import ( + upwind_hflux_miura_cycl_stencil_03a, +) from icon4py.model.common.dimension import EdgeDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -24,9 +25,10 @@ def upwind_hflux_miura_cycl_stencil_03a_numpy( z_tracer_mflx_1_dsl: np.array, z_tracer_mflx_2_dsl: np.array, ): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl)/float(2) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) return p_out_e + def test_upwind_hflux_miura_cycl_stencil_03a(): mesh = SimpleMesh() z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py index a3af39a827..db075b6b39 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import upwind_hflux_miura_cycl_stencil_03b +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import ( + upwind_hflux_miura_cycl_stencil_03b, +) from icon4py.model.common.dimension import EdgeDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -25,11 +26,12 @@ def upwind_hflux_miura_cycl_stencil_03b_numpy( z_tracer_mflx_2_dsl: np.array, z_tracer_mflx_3_dsl: np.array, ): - p_out_e = (z_tracer_mflx_1_dsl - + z_tracer_mflx_2_dsl - + z_tracer_mflx_3_dsl)/float(3) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float( + 3 + ) return p_out_e + def test_upwind_hflux_miura_cycl_stencil_03b(): mesh = SimpleMesh() z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py index 5d65379c80..8ff9f7f03e 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -12,14 +12,17 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.upwind_hflux_miura_stencil_01 import ( upwind_hflux_miura_stencil_01, ) -from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - -from icon4py.model.common.test_utils.helpers import random_field, zero_field, constant_field +from icon4py.model.common.test_utils.helpers import ( + constant_field, + random_field, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -38,9 +41,24 @@ def upwind_hflux_miura_stencil_01_numpy( z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - p_out_e = ( np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) - + distv_bary_1 *np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) - + distv_bary_2 *np.where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + p_out_e = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) ) * p_mass_flx_e return p_out_e diff --git a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 6321d78c8d..86e49c9711 100644 --- a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -17,7 +17,6 @@ upwind_vflux_ppm_stencil_01, ) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py index f9b2586bfd..74da8a299f 100644 --- a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py @@ -13,9 +13,10 @@ import numpy as np -from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 +from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import ( + vert_adv_stencil_01, +) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py index 163eeb900d..b9816e140e 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -18,7 +18,6 @@ v_limit_prbl_sm_stencil_01, ) from icon4py.model.common.dimension import CellDim, KDim - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py index 670664f1f0..6f77e35993 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -18,8 +18,11 @@ v_limit_prbl_sm_stencil_02, ) from icon4py.model.common.dimension import CellDim, KDim - -from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.helpers import ( + random_field, + random_mask, + zero_field, +) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py index 333c21116d..2d7ee1b965 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py @@ -71,7 +71,9 @@ def _apply_diffusion_to_vn( vn, nudgezone_diff, ), - _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v), + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), ) if limited_area else where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 4e8e1585f3..338284be22 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -52,7 +52,9 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdx, dwdy = where( int32(0) < vert_idx, - _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), + _calculate_horizontal_gradients_for_turbulence( + w_old, geofac_grg_x, geofac_grg_y + ), (dwdx, dwdy), ) @@ -69,7 +71,9 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( & (vert_idx < nrdmax) & (interior_idx <= horz_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), + _apply_nabla2_to_w_in_upper_damping_layer( + w, diff_multfac_n2w, area, z_nabla2_c + ), w, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py index 313e9ac500..aa76bed1b8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py @@ -27,7 +27,9 @@ def _apply_nabla2_and_nabla4_global_to_vn( diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + vn = vn + area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge + ) return vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py index c2f1074899..1710922bd0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py @@ -35,4 +35,6 @@ def apply_nabla2_to_vn_in_lateral_boundary( vn: Field[[EdgeDim, KDim], float], fac_bdydiff_v: float, ): - _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn) + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py index 42128fd8ad..3483966bad 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py @@ -36,4 +36,6 @@ def apply_nabla2_to_w_in_upper_damping_layer( cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _apply_nabla2_to_w_in_upper_damping_layer( + w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py index 1c686fe580..60c3126a24 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py @@ -38,4 +38,6 @@ def calculate_diagnostics_for_turbulence( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:])) + _calculate_diagnostics_for_turbulence( + div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:]) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 49eec5e7cf..34c02a07d8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -33,7 +33,9 @@ def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( theta_v, theta_ref_mc, thresh_tdiff ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d + ) return kh_smag_e diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py index 5a0a98d830..d0fff9756b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py @@ -37,4 +37,6 @@ def calculate_horizontal_gradients_for_turbulence( dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], ): - _calculate_horizontal_gradients_for_turbulence(w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy)) + _calculate_horizontal_gradients_for_turbulence( + w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py index 92bf9e30e7..82962fba7a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py @@ -43,4 +43,6 @@ def calculate_nabla2_for_theta( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _calculate_nabla2_for_theta(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _calculate_nabla2_for_theta( + kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 37430b5448..ca00e06c58 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -32,4 +32,6 @@ def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: Field[[EdgeDim, KDim], float], enh_diffu_3d: Field[[CellDim, KDim], float], ): - _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e) + _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d, out=kh_smag_e + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py index d2e875b515..fac30e51fd 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py @@ -45,11 +45,13 @@ def _mo_advection_traj_btraj_compute_o1_dsl( p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( @@ -102,9 +104,11 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=(p_cell_idx, - p_cell_rel_idx_dsl, - p_cell_blk, - p_distv_bary_1, - p_distv_bary_2), + out=( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_distv_bary_1, + p_distv_bary_2, + ), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 295be5cab3..06a4125d86 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -33,4 +33,6 @@ def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( c_intp: Field[[VertexDim, V2CDim], float], p_vert_out: Field[[VertexDim, KDim], float], ): - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(p_cell_in, c_intp, out=p_vert_out) + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( + p_cell_in, c_intp, out=p_vert_out + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py index 0a9bf2d689..170414b2e6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -37,4 +37,6 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], ): - _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) + _mo_intp_rbf_rbf_vec_interpol_vertex( + p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py index ca293f5a91..509a9e6150 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py @@ -24,7 +24,9 @@ def _mo_solve_nonhydro_stencil_02( exner_ref_mc: Field[[CellDim, KDim], float], exner_pr: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_exner_ex_pr = (1.0 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + z_exner_ex_pr = (1.0 + exner_exfac) * ( + exner - exner_ref_mc + ) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return z_exner_ex_pr, exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index b709120573..bfb6b25f66 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -43,7 +43,9 @@ def _mo_solve_nonhydro_stencil_10( ]: z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now(Koff[-1]) + wgt_nnew_rth * rho_var(Koff[-1]) - z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var(Koff[-1]) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var( + Koff[-1] + ) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index eb9ac90b69..a2205e3cb7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -40,11 +40,13 @@ def _compute_btraj( lvn_pos = where(p_vn > 0.0, True, False) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py index 4c628b66cf..f7841c6d0b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py @@ -28,7 +28,10 @@ def _mo_solve_nonhydro_stencil_17( ) -> Field[[EdgeDim, KDim], float]: scalfac_dd3d = broadcast(scalfac_dd3d, (EdgeDim, KDim)) z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) + hmask_dd3d + * scalfac_dd3d + * inv_dual_edge_length + * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) ) return z_graddiv_vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py index 5046f22bfa..9d99e03bab 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py @@ -23,7 +23,9 @@ def _mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = inv_dual_edge_length * (z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0])) + z_gradh_exner = inv_dual_edge_length * ( + z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0]) + ) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py index 0477ff1fcb..453dfec9cc 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py @@ -49,12 +49,18 @@ def _mo_solve_nonhydro_stencil_20( ( 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_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_dexner_dz_c1_0(E2C[0]) + + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) + ) ) ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py index 84fb050f6e..f1919b1934 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py @@ -25,7 +25,9 @@ def _mo_solve_nonhydro_stencil_22( z_hydro_corr: Field[[EdgeDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = where(ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner) + z_gradh_exner = where( + ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner + ) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py index a23544d444..66f3c22346 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py @@ -28,7 +28,9 @@ def _mo_solve_nonhydro_stencil_24( dtime: float, cpd: float, ) -> Field[[EdgeDim, KDim], float]: - vn_nnew = vn_nnow + dtime * (ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner) + vn_nnew = vn_nnow + dtime * ( + ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner + ) return vn_nnew diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index 84abeee0e7..de3b9253b4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -32,7 +32,9 @@ def _mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: z_w_concorr_me_offset_1 = z_w_concorr_me(Koff[-1]) - z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) + z_w_concorr_mc_m1 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim + ) z_w_concorr_mc_m0 = neighbor_sum(e_bln_c_s * z_w_concorr_me(C2E), axis=C2EDim) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 6c40169603..a0adf8bb90 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -35,9 +35,15 @@ def _mo_solve_nonhydro_stencil_40( z_w_concorr_me_offset_2 = z_w_concorr_me(Koff[-2]) z_w_concorr_me_offset_3 = z_w_concorr_me(Koff[-3]) - z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) - z_w_concorr_mc_m2 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim) - z_w_concorr_mc_m3 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim) + z_w_concorr_mc_m1 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim + ) + z_w_concorr_mc_m2 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim + ) + z_w_concorr_mc_m3 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim + ) return ( wgtfacq_c(Koff[-1]) * z_w_concorr_mc_m1 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py index 5b6b0d51e2..8606aa7198 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py @@ -33,7 +33,9 @@ def _mo_solve_nonhydro_stencil_42( cpd: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + + wgt_nnew_vel * ddt_w_adv_ntl2 + - cpd * z_th_ddz_exner_c ) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) return z_w_expl, z_contr_w_fl_l diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py index a354730e7d..d55620484a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py @@ -42,8 +42,17 @@ def _mo_solve_nonhydro_stencil_55( rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( rho_ic * w - rho_ic(Koff[1]) * w(Koff[1]) ) - exner_new = z_exner_expl + exner_ref_mc - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) - theta_v_new = rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new + exner_new = ( + z_exner_expl + + exner_ref_mc + - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) + ) + theta_v_new = ( + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new + ) return rho_new, exner_new, theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py index e5e70e7ba2..b1f718731a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py @@ -24,7 +24,9 @@ def _mo_solve_nonhydro_stencil_56_63( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - z_dwdz_dd = inv_ddqz_z_full * ((w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1]))) + z_dwdz_dd = inv_ddqz_z_full * ( + (w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1])) + ) return z_dwdz_dd diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py index 17879f2dbe..8b5347ada1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py @@ -27,7 +27,9 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], float], r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: - mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) + mass_flx_ic = mass_flx_ic + ( + r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py index 995322d72d..d0435fdde2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py @@ -26,7 +26,9 @@ def _mo_solve_nonhydro_stencil_60( ndyn_substeps_var: float, dtime: float, ) -> Field[[CellDim, KDim], float]: - exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) + exner_dyn_incr = exner - ( + exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy + ) return exner_dyn_incr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py index eea826badc..3ff46a988b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py @@ -30,7 +30,9 @@ def _mo_solve_nonhydro_stencil_65( r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps + * rho_ic + * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py index 6f46f314d2..0c949bd699 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py @@ -39,4 +39,6 @@ def mo_solve_nonhydro_stencil_67( rd_o_cvd: float, rd_o_p0ref: float, ): - _mo_solve_nonhydro_stencil_67(rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner)) + _mo_solve_nonhydro_stencil_67( + rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py index 15230b8afb..8725c196c5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py @@ -31,7 +31,10 @@ def _mo_solve_nonhydro_stencil_68( ) -> Field[[CellDim, KDim], float]: theta_v_new = where( mask_prog_halo_c, - rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new, + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new, theta_v_new, ) return theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py index c7c051da02..72ad7e5471 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py @@ -37,4 +37,6 @@ def mo_velocity_advection_stencil_04( vt: Field[[EdgeDim, KDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me) + _mo_velocity_advection_stencil_04( + vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py index a08de19b3f..4b965b6e8f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py @@ -37,7 +37,9 @@ def _mo_velocity_advection_stencil_07( ) -> Field[[EdgeDim, KDim], float]: return vn_ie * inv_dual_edge_length * ( w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * (z_w_v(E2V[0]) - z_w_v(E2V[1])) + ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * ( + z_w_v(E2V[0]) - z_w_v(E2V[1]) + ) @program(grid_type=GridType.UNSTRUCTURED) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 4dd324603b..4cdbc7d731 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -23,7 +23,9 @@ def _mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc(Koff[-1]) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc( + Koff[-1] + ) return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index 3bce7d7ef2..f413dceaea 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py @@ -43,9 +43,13 @@ def _mo_velocity_advection_stencil_14( vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c) + z_w_con_c = where( + (cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c + ) - z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c) + z_w_con_c = where( + (cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c + ) return cfl_clipping, pre_levelmask, vcfl, z_w_con_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py index a1823ac6be..38b60dd96d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py @@ -26,7 +26,9 @@ def _mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: ddt_w_adv = -z_w_con_c * ( - w(Koff[-1]) * coeff1_dwdz - w(Koff[1]) * coeff2_dwdz + w * (coeff2_dwdz - coeff1_dwdz) + w(Koff[-1]) * coeff1_dwdz + - w(Koff[1]) * coeff2_dwdz + + w * (coeff2_dwdz - coeff1_dwdz) ) return ddt_w_adv @@ -39,4 +41,6 @@ def mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ddt_w_adv: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:]) + _mo_velocity_advection_stencil_16( + z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:] + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py index c4e8d32d2b..9499469aaf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py @@ -45,7 +45,8 @@ def _mo_velocity_advection_stencil_18( ddt_w_adv = where( levelmask & cfl_clipping & owner_mask, - ddt_w_adv + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), + ddt_w_adv + + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), ddt_w_adv, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py index 621e6befed..6476e2a57b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py @@ -45,7 +45,10 @@ def _mo_velocity_advection_stencil_19( ): ddt_vn_adv = -( (coeff_gradekin(E2EC[0]) - coeff_gradekin(E2EC[1])) * z_kin_hor_e - + (-coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1])) + + ( + -coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1]) + ) + vt * (f_e + 0.5 * neighbor_sum(zeta(E2V), axis=E2VDim)) + neighbor_sum(z_w_con_c_full(E2C) * c_lin_e, axis=E2CDim) * (vn_ie - vn_ie(Koff[1])) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py index 173bd2e9b1..abcc04c015 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py @@ -76,7 +76,9 @@ def _mo_velocity_advection_stencil_20( (levelmask | levelmask(Koff[1])) & (abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), ddt_vn_adv + difcoef * area_edge * neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim) - + tangent_orientation * inv_primal_edge_length * neighbor_sum(zeta(E2V), axis=E2VDim), + + tangent_orientation + * inv_primal_edge_length + * neighbor_sum(zeta(E2V), axis=E2VDim), ddt_vn_adv, ) return ddt_vn_adv diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index e331a620cd..262880064a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -46,11 +46,20 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( 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])) + * ( + 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])) + * ( + 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])) + * ( + vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2]) + ) ) z_temp = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py index bc53f68979..e495537311 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py @@ -40,4 +40,6 @@ def update_theta_and_exner( exner: Field[[CellDim, KDim], float], rd_o_cvd: float, ): - _update_theta_and_exner(z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner)) + _update_theta_and_exner( + z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner) + ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py index d9d058adf3..18bc252014 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -44,8 +44,12 @@ def input_data(self, mesh): ) @staticmethod - def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): + def reference( + mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn + ): area_edge = np.expand_dims(area_edge, axis=-1) diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + vn = vn + area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge + ) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py index c3f4725214..42af67e615 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py @@ -34,7 +34,9 @@ def reference( ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum(z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1) + w = w - diff_multfac_w * area * area * np.sum( + z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1 + ) return dict(w=w) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py index c60af2834b..71da351f7d 100644 --- a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py @@ -36,7 +36,9 @@ def reference( kc_offset_1 = np.roll(kh_c, shift=1, axis=1) div_offset_1 = np.roll(div, shift=1, axis=1) div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] - hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[ + :, 1: + ] return dict(div_ic=div_ic, hdef_ic=hdef_ic) @pytest.fixture @@ -46,4 +48,6 @@ def input_data(self, mesh): kh_c = random_field(mesh, CellDim, KDim) div_ic = zero_field(mesh, CellDim, KDim) hdef_ic = zero_field(mesh, CellDim, KDim) - return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) + return dict( + wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic + ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py index 0d6cb63d3a..4fed66fbaa 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py @@ -43,7 +43,9 @@ def test_calculate_nabla2_of_theta(): out = zero_field(mesh, CellDim, KDim) - ref = calculate_nabla2_of_theta_numpy(mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div)) + ref = calculate_nabla2_of_theta_numpy( + mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div) + ) calculate_nabla2_of_theta( z_nabla2_e, geofac_div_new, diff --git a/model/atmosphere/dycore/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py index e67739f53f..307dbdaf25 100644 --- a/model/atmosphere/dycore/tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/test_compute_airmass.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np - from icon4py.atm_dyn_iconam.compute_airmass import compute_airmass from icon4py.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index a38c366a5c..dff37381e3 100644 --- a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -54,22 +54,28 @@ def mo_advection_traj_btraj_compute_o1_dsl_numpy( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py index 4238b31fd3..0d5c38403b 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py @@ -38,7 +38,9 @@ def reference( exner_exfac: np.array, **kwargs, ) -> dict: - z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + z_exner_ex_pr = (1 + exner_exfac) * ( + exner - exner_ref_mc + ) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py index 29cb23efed..04df89bb93 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py @@ -37,9 +37,12 @@ def reference( z_exner_ic: np.array, ) -> np.array: z_exner_ic[:, 3:] = ( - np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) - + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) - + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) + np.roll(wgtfacq_c, shift=1, axis=1) + * np.roll(z_exner_ex_pr, shift=1, axis=1) + + np.roll(wgtfacq_c, shift=2, axis=1) + * np.roll(z_exner_ex_pr, shift=2, axis=1) + + np.roll(wgtfacq_c, shift=3, axis=1) + * np.roll(z_exner_ex_pr, shift=3, axis=1) )[:, 3:] return {"z_exner_ic": z_exner_ic} diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py index 9a86409fa4..2ba664c629 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil06(StencilTest): OUTPUTS = ("z_dexner_dz_c_1",) @staticmethod - def reference(mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs) -> np.array: + def reference( + mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs + ) -> np.array: z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full return dict(z_dexner_dz_c_1=z_dexner_dz_c_1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py index bb92f1fbef..40f10d2032 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil15(StencilTest): OUTPUTS = ("z_rho_e", "z_theta_v_e") @staticmethod - def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> tuple[np.array]: + def reference( + mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs + ) -> tuple[np.array]: z_rho_e = np.zeros_like(z_rho_e) z_theta_v_e = np.zeros_like(z_theta_v_e) return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 8282ea91e1..4e5e50ed28 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -45,22 +45,28 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_distv_bary_1, p_distv_bary_2 @@ -194,7 +200,10 @@ def test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1(): z_rho_e = random_field(mesh, EdgeDim, KDim) z_theta_v_e = random_field(mesh, EdgeDim, KDim) - (z_rho_e_ref, z_theta_v_e_ref,) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( + ( + z_rho_e_ref, + z_theta_v_e_ref, + ) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( mesh.e2c, np.asarray(p_vn), np.asarray(p_vt), diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py index 5a782f7c5f..56d05c7b98 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py @@ -41,8 +41,9 @@ def reference( z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( - c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1 + z_gradh_exner = ( + inv_dual_edge_length * z_exner_ex_weighted + - ddxn_z_full * np.sum(c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1) ) return dict(z_gradh_exner=z_gradh_exner) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py index 349355cd72..8e5bc613be 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py @@ -51,8 +51,12 @@ 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, :] * ( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py index e0f31fb5ef..4423ac168d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py @@ -61,7 +61,9 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): full_shape, theta_v_ic, e2c, ikoffset ) - inv_ddqz_z_full_at_kidx, _ = _apply_index_field(full_shape, inv_ddqz_z_full, e2c, ikoffset) + inv_ddqz_z_full_at_kidx, _ = _apply_index_field( + full_shape, inv_ddqz_z_full, e2c, ikoffset + ) z_theta1 = ( theta_v_at_kidx[:, 0, :] diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py index 18623f00db..cac9fbeccd 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil25(StencilTest): OUTPUTS = ("z_graddiv2_vn",) @staticmethod - def reference(mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs) -> np.array: + def reference( + mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs + ) -> np.array: geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) z_graddiv2_vn = np.sum(z_graddiv_vn[mesh.e2c2eO] * geofac_grdiv, axis=1) return dict(z_graddiv2_vn=z_graddiv2_vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py index ec2389a8de..dded703d10 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil26(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference(mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs) -> dict: + def reference( + mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs + ) -> dict: vn = vn + (scal_divdamp_o2 * z_graddiv_vn) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py index 17234d448c..2eeba98f41 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil28(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference(mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs) -> np.array: + def reference( + mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs + ) -> np.array: vn = vn + (iau_wgt_dyn * vn_incr) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py index da69d23da3..deb5bf35cc 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil29(StencilTest): OUTPUTS = ("vn_new",) @staticmethod - def reference(mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs) -> dict: + def reference( + mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs + ) -> dict: vn_new = vn_now + dtime * grf_tend_vn return dict(vn_new=vn_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py index 7e0ea0bdd9..7af31248f2 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py @@ -46,7 +46,9 @@ def reference( **kwargs, ) -> tuple[np.array]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + + wgt_nnew_vel * ddt_w_adv_ntl2 + - cpd * z_th_ddz_exner_c ) vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py index 2660d320ef..6624930d7b 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py @@ -44,7 +44,9 @@ def reference( cvd, **kwargs, ) -> dict: - z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full + z_beta = ( + dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full + ) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py index 60253bc78e..3f477369d1 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil47(StencilTest): OUTPUTS = ("w_nnew", "z_contr_w_fl_l") @staticmethod - def reference(mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: + def reference( + mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs + ) -> dict: w_nnew = w_concorr_c z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py index f7e12c96f2..d1cb07a661 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py @@ -57,7 +57,9 @@ def mo_solve_nonhydro_stencil_51_numpy( z_c = -z_gamma * z_beta * z_alpha_k_plus_1 z_b = 1.0 + z_gamma * z_alpha[:, :-1] * (z_beta_k_minus_1 + z_beta) z_q = mo_solve_nonhydro_stencil_51_z_q_numpy(z_c, z_b) - w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy(z_gamma, z_b, z_w_expl, z_exner_expl) + w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy( + z_gamma, z_b, z_w_expl, z_exner_expl + ) return z_q, w_nnew diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py index 6c0cc28e0d..85b1531896 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py @@ -49,11 +49,15 @@ def mo_solve_nonhydro_stencil_52_numpy( for k in range(1, k_size): z_a[:, k] = -z_gamma[:, k] * z_beta[:, k - 1] * z_alpha[:, k - 1] z_c[:, k] = -z_gamma[:, k] * z_beta[:, k] * z_alpha[:, k + 1] - z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * (z_beta[:, k - 1] + z_beta[:, k]) + z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * ( + z_beta[:, k - 1] + z_beta[:, k] + ) z_g[:, k] = 1.0 / (z_b[:, k] + z_a[:, k] * z_q[:, k - 1]) z_q[:, k] = -z_c[:, k] * z_g[:, k] - w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * (z_exner_expl[:, k - 1] - z_exner_expl[:, k]) + w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * ( + z_exner_expl[:, k - 1] - z_exner_expl[:, k] + ) w[:, k] = (w[:, k] - z_a[:, k] * w[:, k - 1]) * z_g[:, k] return z_q, w diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py index eeb17c2cce..bc2e501a21 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil54(StencilTest): OUTPUTS = ("w",) @staticmethod - def reference(mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -> np.array: + def reference( + mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs + ) -> np.array: z_raylfac = np.expand_dims(z_raylfac, axis=0) w_1 = np.expand_dims(w_1, axis=-1) w = z_raylfac * w + (1.0 - z_raylfac) * w_1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py index 1d9e126493..cd6a2683c4 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py @@ -62,7 +62,10 @@ def reference( - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) ) theta_v_new = ( - rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new ) return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py index dff001052b..93e8c6a95d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py @@ -37,7 +37,9 @@ def reference( **kwargs, ) -> dict: vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) + mass_flx_ic = mass_flx_ic + ( + r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + ) return dict(mass_flx_ic=mass_flx_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py index 634446cb33..407dc06e90 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py @@ -35,7 +35,9 @@ def reference( dtime: float, **kwargs, ) -> np.array: - exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) + exner_dyn_incr = exner - ( + exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy + ) return dict(exner_dyn_incr=exner_dyn_incr) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py index 41ef7c9439..8ff168cd3d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil62(StencilTest): OUTPUTS = ("w_new",) @staticmethod - def reference(mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs) -> np.array: + def reference( + mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs + ) -> np.array: w_new = w_now + dtime * grf_tend_w return dict(w_new=w_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py index b5b5b9c0c0..77e46f5927 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py @@ -41,7 +41,9 @@ def reference( vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps + * rho_ic + * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return dict(mass_flx_ic=mass_flx_ic) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py index dd54adca2b..297ec945da 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py @@ -46,7 +46,10 @@ def reference( theta_v_new = np.where( mask_prog_halo_c, - rho_now * theta_v_now * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) / rho_new, + rho_now + * theta_v_now + * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) + / rho_new, theta_v_new, ) return dict(theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index 9bdc11bf27..aa64258d76 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -21,13 +21,17 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: +def mo_velocity_advection_stencil_02_vn_ie_numpy( + wgtfac_e: np.array, vn: np.array +) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie -def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: +def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( + vn: np.array, vt: np.array +) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py index f32c87b1b1..d7c3f41893 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py @@ -21,7 +21,9 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_03_numpy(wgtfac_e: np.array, vt: np.array) -> np.array: +def mo_velocity_advection_stencil_03_numpy( + wgtfac_e: np.array, vt: np.array +) -> np.array: vt_k_minus_1 = np.roll(vt, shift=1, axis=1) z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 @@ -36,7 +38,9 @@ def test_mo_velocity_advection_stencil_03(): z_vt_ie = zero_field(mesh, EdgeDim, KDim) - z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy(np.asarray(wgtfac_e), np.asarray(vt)) + z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy( + np.asarray(wgtfac_e), np.asarray(vt) + ) mo_velocity_advection_stencil_03( wgtfac_e, vt, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py index 1acd23aeb1..f1eefbb0c6 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py @@ -21,7 +21,9 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> np.array: +def mo_velocity_advection_stencil_06_numpy( + wgtfacq_e: np.array, vn: np.array +) -> np.array: vn_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_k_minus_2 = np.roll(vn, shift=2, axis=1) vn_k_minus_3 = np.roll(vn, shift=3, axis=1) @@ -45,7 +47,9 @@ def test_mo_velocity_advection_stencil_06(): vn_ie = zero_field(mesh, EdgeDim, KDim) - vn_ie_ref = mo_velocity_advection_stencil_06_numpy(np.asarray(wgtfacq_e), np.asarray(vn)) + vn_ie_ref = mo_velocity_advection_stencil_06_numpy( + np.asarray(wgtfacq_e), np.asarray(vn) + ) mo_velocity_advection_stencil_06( wgtfacq_e, vn, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py index 192de3f3c3..80b8a8b936 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py @@ -30,7 +30,9 @@ class TestMoVelocityAdvectionStencil08(StencilTest): OUTPUTS = ("z_ekinh",) @staticmethod - def reference(mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + def reference( + mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs + ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_ekinh = np.sum(z_kin_hor_e[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_ekinh=z_ekinh) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py index 664b168561..e02b7c37ed 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py @@ -30,7 +30,9 @@ class TestMoVelocityAdvectionStencil09(StencilTest): OUTPUTS = ("z_w_concorr_mc",) @staticmethod - def reference(mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + def reference( + mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs + ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_mc = np.sum(z_w_concorr_me[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_w_concorr_mc=z_w_concorr_mc) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index 33c9bfe847..19f5e0f325 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -25,7 +25,9 @@ def mo_velocity_advection_stencil_10_numpy( wgtfac_c: np.array, z_w_concorr_mc: np.array ) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + w_concorr_c = ( + wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + ) return w_concorr_c diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py index eb7fb3fc44..336ce24d35 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py @@ -26,7 +26,9 @@ class TestMoVelocityAdvectionStencil13(StencilTest): OUTPUTS = ("z_w_con_c",) @staticmethod - def reference(mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> np.array: + def reference( + mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs + ) -> np.array: z_w_con_c = z_w_con_c - w_concorr_c return dict(z_w_con_c=z_w_con_c) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py index eaa6532beb..48d9cb6b37 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py @@ -66,7 +66,10 @@ def mo_velocity_advection_stencil_19_numpy( ddt_vn_adv = -( (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + + ( + -coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1] + ) + vt * (f_e + 0.5 * np.sum(zeta[e2v], axis=1)) + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) * (vn_ie[:, :-1] - vn_ie[:, 1:]) diff --git a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 123a8ab301..8c2209a134 100644 --- a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -56,12 +56,15 @@ def mo_nh_diffusion_stencil_15_numpy( ] sum_over = np.sum( - geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + geofac_n2s_nbh + * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), axis=1, ) geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - return np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) + return np.where( + mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp + ) def test_mo_nh_diffusion_stencil_15(): From ab1bf91643b908c815a2566e6052a6c227df7ef0 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 24 Aug 2023 17:42:43 +0200 Subject: [PATCH 087/105] Place compute_airmass.py and btraj_01.py in right directory and adapt import path --- .../advection}/mo_advection_traj_btraj_compute_o1_dsl.py | 0 .../tests/test_mo_advection_traj_btraj_compute_o1_dsl.py | 9 +++++++-- model/atmosphere/dycore/tests/test_compute_airmass.py | 8 ++++---- 3 files changed, 11 insertions(+), 6 deletions(-) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => advection/src/icon4py/model/atmosphere/advection}/mo_advection_traj_btraj_compute_o1_dsl.py (100%) rename model/atmosphere/{dycore => advection}/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py (92%) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py diff --git a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 92% rename from model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index dff37381e3..c366735626 100644 --- a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.mo_advection_traj_btraj_compute_o1_dsl import ( +from icon4py.model.atmosphere.advection.mo_advection_traj_btraj_compute_o1_dsl import ( mo_advection_traj_btraj_compute_o1_dsl, ) from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim @@ -51,6 +51,7 @@ def mo_advection_traj_btraj_compute_o1_dsl_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, np.int32(0), np.int32(1)) p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( @@ -78,7 +79,7 @@ def mo_advection_traj_btraj_compute_o1_dsl_numpy( + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) - return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 + return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 def test_mo_advection_traj_btraj_compute_o1_dsl(): @@ -103,6 +104,7 @@ def test_mo_advection_traj_btraj_compute_o1_dsl(): dual_normal_cell_2 = random_field(mesh, EdgeDim, E2CDim) dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) p_cell_idx = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_cell_blk = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_distv_bary_1 = random_field(mesh, EdgeDim, KDim) p_distv_bary_2 = random_field(mesh, EdgeDim, KDim) @@ -110,6 +112,7 @@ def test_mo_advection_traj_btraj_compute_o1_dsl(): ( p_cell_idx_ref, + p_cell_rel_idx_dsl_ref, p_cell_blk_ref, p_distv_bary_1_ref, p_distv_bary_2_ref, @@ -139,6 +142,7 @@ def test_mo_advection_traj_btraj_compute_o1_dsl(): primal_normal_cell_2_new, dual_normal_cell_2_new, p_cell_idx, + p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2, @@ -149,6 +153,7 @@ def test_mo_advection_traj_btraj_compute_o1_dsl(): }, ) assert np.allclose(p_cell_idx, p_cell_idx_ref) + assert np.allclose(p_cell_rel_idx_dsl, p_cell_rel_idx_dsl_ref) assert np.allclose(p_cell_blk, p_cell_blk_ref) assert np.allclose(p_distv_bary_1, p_distv_bary_1_ref) assert np.allclose(p_distv_bary_2, p_distv_bary_2_ref) diff --git a/model/atmosphere/dycore/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py index 307dbdaf25..406c755810 100644 --- a/model/atmosphere/dycore/tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/test_compute_airmass.py @@ -12,11 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from icon4py.atm_dyn_iconam.compute_airmass import compute_airmass -from icon4py.common.dimension import CellDim, KDim +from icon4py.model.atmosphere.dycore.compute_airmass import compute_airmass +from icon4py.model.common.dimension import EdgeDim, KDim, CellDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from .test_utils.helpers import random_field -from .test_utils.simple_mesh import SimpleMesh def compute_airmass_numpy( From fdbb38aba7acc2e47ef09b3fd74f36616674bf12 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 25 Aug 2023 15:27:41 +0200 Subject: [PATCH 088/105] Remove __init__.py in advection/tests/ and dycore/tests for pytest to work + adapt tox files --- model/atmosphere/advection/tests/__init__.py | 12 ------------ model/atmosphere/dycore/tests/__init__.py | 12 ------------ model/tox.ini | 3 +-- tox.ini | 3 +-- 4 files changed, 2 insertions(+), 28 deletions(-) delete mode 100644 model/atmosphere/advection/tests/__init__.py delete mode 100644 model/atmosphere/dycore/tests/__init__.py diff --git a/model/atmosphere/advection/tests/__init__.py b/model/atmosphere/advection/tests/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/model/atmosphere/advection/tests/__init__.py +++ /dev/null @@ -1,12 +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 diff --git a/model/atmosphere/dycore/tests/__init__.py b/model/atmosphere/dycore/tests/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/model/atmosphere/dycore/tests/__init__.py +++ /dev/null @@ -1,12 +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 diff --git a/model/tox.ini b/model/tox.ini index b97a24f41e..e14438ea04 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,8 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/advection/src common/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/advection/src common/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html diff --git a/tox.ini b/tox.ini index 4768cc709e..eec85c5901 100644 --- a/tox.ini +++ b/tox.ini @@ -14,8 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/common/src tools/src - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/advection/src model/common/src tools/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/advection/src model/common/src tools/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html From eec80576d9d94a18ad7fb8ee31add699f69083a7 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 25 Aug 2023 15:41:21 +0200 Subject: [PATCH 089/105] Modify recon_lsq_cell_c_stencil.py (not used in ICON code path) --- .../advection/recon_lsq_cell_c_stencil.py | 18 -- .../tests/test_recon_lsq_cell_c_stencil.py | 253 ++++++++---------- 2 files changed, 114 insertions(+), 157 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py index d22089f684..310bff8ef2 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py @@ -107,24 +107,6 @@ def _recon_lsq_cell_c_stencil( Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], ]: - lsq_rmat_rdiag_c_1 = broadcast(lsq_rmat_rdiag_c_1, (CellDim, KDim)) - lsq_rmat_rdiag_c_2 = broadcast(lsq_rmat_rdiag_c_2, (CellDim, KDim)) - lsq_rmat_rdiag_c_3 = broadcast(lsq_rmat_rdiag_c_3, (CellDim, KDim)) - lsq_rmat_rdiag_c_4 = broadcast(lsq_rmat_rdiag_c_4, (CellDim, KDim)) - lsq_rmat_rdiag_c_5 = broadcast(lsq_rmat_rdiag_c_5, (CellDim, KDim)) - lsq_rmat_rdiag_c_6 = broadcast(lsq_rmat_rdiag_c_6, (CellDim, KDim)) - lsq_rmat_rdiag_c_7 = broadcast(lsq_rmat_rdiag_c_7, (CellDim, KDim)) - lsq_rmat_rdiag_c_8 = broadcast(lsq_rmat_rdiag_c_8, (CellDim, KDim)) - lsq_rmat_rdiag_c_9 = broadcast(lsq_rmat_rdiag_c_9, (CellDim, KDim)) - lsq_qtmat_c_1 = broadcast(lsq_qtmat_c_1, (CECECDim, KDim)) - lsq_qtmat_c_2 = broadcast(lsq_qtmat_c_2, (CECECDim, KDim)) - lsq_qtmat_c_3 = broadcast(lsq_qtmat_c_3, (CECECDim, KDim)) - lsq_qtmat_c_4 = broadcast(lsq_qtmat_c_4, (CECECDim, KDim)) - lsq_qtmat_c_5 = broadcast(lsq_qtmat_c_5, (CECECDim, KDim)) - lsq_qtmat_c_6 = broadcast(lsq_qtmat_c_6, (CECECDim, KDim)) - lsq_qtmat_c_7 = broadcast(lsq_qtmat_c_7, (CECECDim, KDim)) - lsq_qtmat_c_8 = broadcast(lsq_qtmat_c_8, (CECECDim, KDim)) - lsq_qtmat_c_9 = broadcast(lsq_qtmat_c_9, (CECECDim, KDim)) p_coeff_10 = lsq_rmat_rdiag_c_9 * ( lsq_qtmat_c_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py index 916e5df693..f7a50a179c 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py @@ -11,6 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import sys # Increase recusion depth, otherwise it doesn't compile + import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider @@ -26,6 +28,9 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh +sys.setrecursionlimit(6000) + + def recon_lsq_cell_c_stencil_numpy( c2e2c2e2c: np.ndarray, p_cc: np.ndarray, @@ -93,27 +98,6 @@ def recon_lsq_cell_c_stencil_numpy( lsq_moments_8: np.ndarray, lsq_moments_9: np.ndarray, ) -> tuple[np.ndarray]: - p_cc_e = np.expand_dims(p_cc, axis=-1) - # n_diff = p_cc[c2e2c] - p_cc_e - - lsq_rmat_rdiag_c_1 = np.expand_dims(lsq_rmat_rdiag_c_1, axis=-1) - lsq_rmat_rdiag_c_2 = np.expand_dims(lsq_rmat_rdiag_c_2, axis=-1) - lsq_rmat_rdiag_c_3 = np.expand_dims(lsq_rmat_rdiag_c_3, axis=-1) - lsq_rmat_rdiag_c_4 = np.expand_dims(lsq_rmat_rdiag_c_4, axis=-1) - lsq_rmat_rdiag_c_5 = np.expand_dims(lsq_rmat_rdiag_c_5, axis=-1) - lsq_rmat_rdiag_c_6 = np.expand_dims(lsq_rmat_rdiag_c_6, axis=-1) - lsq_rmat_rdiag_c_7 = np.expand_dims(lsq_rmat_rdiag_c_7, axis=-1) - lsq_rmat_rdiag_c_8 = np.expand_dims(lsq_rmat_rdiag_c_8, axis=-1) - lsq_rmat_rdiag_c_9 = np.expand_dims(lsq_rmat_rdiag_c_9, axis=-1) - # lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) - # lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) - # lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) - # lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) - # lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) - # lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) - # lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) - # lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) - # lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) @@ -132,6 +116,24 @@ def recon_lsq_cell_c_stencil_numpy( lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) + lsq_rmat_rdiag_c_1 = np.expand_dims(lsq_rmat_rdiag_c_1, axis=-1) + lsq_rmat_rdiag_c_2 = np.expand_dims(lsq_rmat_rdiag_c_2, axis=-1) + lsq_rmat_rdiag_c_3 = np.expand_dims(lsq_rmat_rdiag_c_3, axis=-1) + lsq_rmat_rdiag_c_4 = np.expand_dims(lsq_rmat_rdiag_c_4, axis=-1) + lsq_rmat_rdiag_c_5 = np.expand_dims(lsq_rmat_rdiag_c_5, axis=-1) + lsq_rmat_rdiag_c_6 = np.expand_dims(lsq_rmat_rdiag_c_6, axis=-1) + lsq_rmat_rdiag_c_7 = np.expand_dims(lsq_rmat_rdiag_c_7, axis=-1) + lsq_rmat_rdiag_c_8 = np.expand_dims(lsq_rmat_rdiag_c_8, axis=-1) + lsq_rmat_rdiag_c_9 = np.expand_dims(lsq_rmat_rdiag_c_9, axis=-1) + lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) + lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) + lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) + lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) + lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) + lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) + lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) + lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) + lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) lsq_rmat_utri_c_1 = np.expand_dims(lsq_rmat_utri_c_1, axis=-1) lsq_rmat_utri_c_2 = np.expand_dims(lsq_rmat_utri_c_2, axis=-1) lsq_rmat_utri_c_3 = np.expand_dims(lsq_rmat_utri_c_3, axis=-1) @@ -204,91 +206,64 @@ def recon_lsq_cell_c_stencil_numpy( lsq_rmat_utri_c_34 = np.broadcast_to(lsq_rmat_utri_c_34, p_cc.shape) lsq_rmat_utri_c_35 = np.broadcast_to(lsq_rmat_utri_c_35, p_cc.shape) lsq_rmat_utri_c_36 = np.broadcast_to(lsq_rmat_utri_c_36, p_cc.shape) - # lsq_rmat_utri_c_9 lsq_qtmat_c_1 = nplsq_rmat_utri_c_9.broadcast_to(lsq_qtmat_c_1, (CECECDim, KDim)) - # lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, (CECECDim, KDim)) - # lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, (CECECDim, KDim)) - # lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, (CECECDim, KDim)) - # lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, (CECECDim, KDim)) - # lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, (CECECDim, KDim)) - # lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, (CECECDim, KDim)) - # lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, (CECECDim, KDim)) - # lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, (CECECDim, KDim)) - # lsq_qtmat_c_1 = np.broadcast_to(lsq_qtmat_c_1, p_cc.shape) - # lsq_qtmat_c_2 = np.broadcast_to(lsq_qtmat_c_2, p_cc.shape) - # lsq_qtmat_c_3 = np.broadcast_to(lsq_qtmat_c_3, p_cc.shape) - # lsq_qtmat_c_4 = np.broadcast_to(lsq_qtmat_c_4, p_cc.shape) - # lsq_qtmat_c_5 = np.broadcast_to(lsq_qtmat_c_5, p_cc.shape) - # lsq_qtmat_c_6 = np.broadcast_to(lsq_qtmat_c_6, p_cc.shape) - # lsq_qtmat_c_7 = np.broadcast_to(lsq_qtmat_c_7, p_cc.shape) - # lsq_qtmat_c_8 = np.broadcast_to(lsq_qtmat_c_8, p_cc.shape) - # lsq_qtmat_c_9 = np.broadcast_to(lsq_qtmat_c_9, p_cc.shape) - # lsq_qtmat_c_1 = np.repeat(lsq_qtmat_c_1[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_2 = np.repeat(lsq_qtmat_c_2[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_3 = np.repeat(lsq_qtmat_c_3[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_4 = np.repeat(lsq_qtmat_c_4[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_5 = np.repeat(lsq_qtmat_c_5[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_6 = np.repeat(lsq_qtmat_c_6[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_7 = np.repeat(lsq_qtmat_c_7[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_8 = np.repeat(lsq_qtmat_c_8[:, np.newaxis], p_cc.shape[1], axis=1) - # lsq_qtmat_c_9 = np.repeat(lsq_qtmat_c_9[:, np.newaxis], p_cc.shape[1], axis=1) - lsq_qtmat_c_9 = np.expand_dims(lsq_qtmat_c_9, axis=-1) - lsq_qtmat_c_8 = np.expand_dims(lsq_qtmat_c_8, axis=-1) - lsq_qtmat_c_7 = np.expand_dims(lsq_qtmat_c_7, axis=-1) - lsq_qtmat_c_6 = np.expand_dims(lsq_qtmat_c_6, axis=-1) - lsq_qtmat_c_5 = np.expand_dims(lsq_qtmat_c_5, axis=-1) - lsq_qtmat_c_4 = np.expand_dims(lsq_qtmat_c_4, axis=-1) - lsq_qtmat_c_3 = np.expand_dims(lsq_qtmat_c_3, axis=-1) - lsq_qtmat_c_2 = np.expand_dims(lsq_qtmat_c_2, axis=-1) lsq_qtmat_c_1 = np.expand_dims(lsq_qtmat_c_1, axis=-1) + lsq_qtmat_c_2 = np.expand_dims(lsq_qtmat_c_2, axis=-1) + lsq_qtmat_c_3 = np.expand_dims(lsq_qtmat_c_3, axis=-1) + lsq_qtmat_c_4 = np.expand_dims(lsq_qtmat_c_4, axis=-1) + lsq_qtmat_c_5 = np.expand_dims(lsq_qtmat_c_5, axis=-1) + lsq_qtmat_c_6 = np.expand_dims(lsq_qtmat_c_6, axis=-1) + lsq_qtmat_c_7 = np.expand_dims(lsq_qtmat_c_7, axis=-1) + lsq_qtmat_c_8 = np.expand_dims(lsq_qtmat_c_8, axis=-1) + lsq_qtmat_c_9 = np.expand_dims(lsq_qtmat_c_9, axis=-1) p_coeff_10 = lsq_rmat_rdiag_c_9 * ( - lsq_qtmat_c_9[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_9[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_9[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_9[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_9[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_9[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_9[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_9[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_9[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) ) p_coeff_9 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_8[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_8[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_8[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_8[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_8[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_8[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_8[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_8[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_8[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - lsq_rmat_utri_c_1 * p_coeff_10 ) - p_coeff_8 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_7[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_7[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_7[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_7[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_7[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_7[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_7[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_7[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_7[:, 8] * (p_cc_e[:, 8] - p_cc) + p_coeff_8 = lsq_rmat_rdiag_c_7 * ( + lsq_qtmat_c_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) ) p_coeff_7 = lsq_rmat_rdiag_c_6 * ( - lsq_qtmat_c_6[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_6[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_6[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_6[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_6[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_6[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_6[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_6[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_6[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_4 * p_coeff_8 + lsq_rmat_utri_c_5 * p_coeff_9 @@ -297,15 +272,15 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_6 = lsq_rmat_rdiag_c_5 * ( - lsq_qtmat_c_5[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_5[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_5[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_5[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_5[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_5[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_5[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_5[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_5[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_7 * p_coeff_7 + lsq_rmat_utri_c_8 * p_coeff_8 @@ -315,15 +290,15 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_5 = lsq_rmat_rdiag_c_4 * ( - lsq_qtmat_c_4[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_4[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_4[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_4[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_4[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_4[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_4[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_4[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_4[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_11 * p_coeff_6 + lsq_rmat_utri_c_12 * p_coeff_7 @@ -334,15 +309,15 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_4 = lsq_rmat_rdiag_c_3 * ( - lsq_qtmat_c_3[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_3[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_3[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_3[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_3[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_3[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_3[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_3[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_3[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_16 * p_coeff_5 + lsq_rmat_utri_c_17 * p_coeff_6 @@ -354,15 +329,15 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_3 = lsq_rmat_rdiag_c_2 * ( - lsq_qtmat_c_2[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_2[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_2[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_2[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_2[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_2[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_2[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_2[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_2[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_22 * p_coeff_4 + lsq_rmat_utri_c_23 * p_coeff_5 @@ -375,15 +350,15 @@ def recon_lsq_cell_c_stencil_numpy( ) p_coeff_2 = lsq_rmat_rdiag_c_1 * ( - lsq_qtmat_c_1[:, 0] * (p_cc_e[:, 0] - p_cc) - + lsq_qtmat_c_1[:, 1] * (p_cc_e[:, 1] - p_cc) - + lsq_qtmat_c_1[:, 2] * (p_cc_e[:, 2] - p_cc) - + lsq_qtmat_c_1[:, 3] * (p_cc_e[:, 3] - p_cc) - + lsq_qtmat_c_1[:, 4] * (p_cc_e[:, 4] - p_cc) - + lsq_qtmat_c_1[:, 5] * (p_cc_e[:, 5] - p_cc) - + lsq_qtmat_c_1[:, 6] * (p_cc_e[:, 6] - p_cc) - + lsq_qtmat_c_1[:, 7] * (p_cc_e[:, 7] - p_cc) - + lsq_qtmat_c_1[:, 8] * (p_cc_e[:, 8] - p_cc) + lsq_qtmat_c_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_qtmat_c_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_qtmat_c_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_qtmat_c_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_qtmat_c_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_qtmat_c_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_qtmat_c_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_qtmat_c_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_qtmat_c_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ( lsq_rmat_utri_c_29 * p_coeff_3 + lsq_rmat_utri_c_30 * p_coeff_4 @@ -662,7 +637,7 @@ def test_recon_lsq_cell_c_stencil(): p_coeff_9, p_coeff_10, offset_provider={ - "C2E2C2E2C": mesh.get_c2e2c_offset_provider(), + "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), "C2CECEC": StridedNeighborOffsetProvider( CellDim, CECECDim, mesh.n_c2e2c2e2c ), From d85b366fce313079b73bf5611d99ac80725aa227 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 25 Aug 2023 16:18:29 +0200 Subject: [PATCH 090/105] Correct for passing flake8 code style checks --- .../advection/btraj_dreg_stencil_03.py | 6 ++---- .../advection/recon_lsq_cell_c_stencil.py | 1 - .../upwind_hflux_miura_cycl_stencil_02.py | 8 +------- .../upwind_hflux_miura_cycl_stencil_03a.py | 2 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 1 - .../tests/test_btraj_dreg_stencil_01.py | 12 +++--------- .../tests/test_btraj_dreg_stencil_02.py | 3 +-- .../tests/test_btraj_dreg_stencil_03.py | 5 +---- .../test_divide_flux_area_list_stencil_01.py | 2 +- .../test_divide_flux_area_list_stencil_02.py | 3 +-- .../test_hflux_ffsl_hybrid_stencil_01a.py | 2 -- .../tests/test_hflux_ffsl_hybrid_stencil_02.py | 9 +-------- .../tests/test_hflx_limiter_mo_stencil_01a.py | 9 +-------- .../tests/test_hflx_limiter_mo_stencil_01b.py | 5 ----- ...est_prep_gauss_quadrature_c_list_stencil.py | 4 +--- .../test_prep_gauss_quadrature_c_stencil.py | 9 ++------- .../tests/test_recon_lsq_cell_c_svd_stencil.py | 7 +++---- .../test_upwind_hflux_miura3_stencil_01.py | 1 - .../test_upwind_hflux_miura_cycl_stencil_01.py | 1 - .../test_upwind_hflux_miura_cycl_stencil_02.py | 8 +------- .../dycore/tests/test_compute_airmass.py | 6 +++--- .../model/common/test_utils/simple_mesh.py | 18 +++++++++--------- 22 files changed, 32 insertions(+), 90 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index bc120fe1b9..8235d1b611 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -1,5 +1,3 @@ -# TODO: This license is not consistent with license used in the project. -# Delete the inconsistent license and above line and rerun pre-commit to insert a good license. # ICON4Py - ICON inspired code in Python and GT4Py # # Copyright (c) 2022, ETH Zurich and MeteoSwiss @@ -7,11 +5,11 @@ # # 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 lat_dsler +# 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-lat_dsler +# SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py index 310bff8ef2..abea9c0f40 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py @@ -15,7 +15,6 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import broadcast from icon4py.model.common.dimension import ( C2CECEC, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index cb7c1809af..c4ae3fa8ec 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -12,13 +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, - broadcast, - int32, - neighbor_sum, - where, -) +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py index 7c27ef4ee9..a8bcf8117f 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -13,7 +13,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import float64, int32 +from gt4py.next.ffront.fbuiltins import float64 from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py index d127c50224..4f5f73ac91 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -13,7 +13,6 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py index 4e4a3cabd7..b136704f95 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py @@ -12,19 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import ( btraj_dreg_stencil_01, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - _shape, - random_field, - zero_field, -) -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh, zero_field def btraj_dreg_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py index d90c063df5..ecd96ed72a 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py @@ -18,9 +18,8 @@ from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import ( btraj_dreg_stencil_02, ) -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import ECDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - _shape, as_1D_sparse_field, random_field, zero_field, diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py index 42265f5486..4c3a0e088a 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py @@ -13,19 +13,16 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import ( btraj_dreg_stencil_03, ) -from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - _shape, as_1D_sparse_field, constant_field, random_field, - zero_field, ) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py index 9f11ad4f54..0402167fa0 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py @@ -18,7 +18,7 @@ from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( divide_flux_area_list_stencil_01, ) -from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( as_1D_sparse_field, random_field, diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py index e920f1d738..fceb4234db 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py @@ -18,12 +18,11 @@ from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( divide_flux_area_list_stencil_02, ) -from icon4py.model.common.dimension import CellDim, E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( as_1D_sparse_field, random_field, random_mask, - zero_field, ) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py index 024dd0ae90..826c0fdfe1 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -13,14 +13,12 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( hflux_ffsl_hybrid_stencil_01a, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - _shape, constant_field, random_field, zero_field, diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py index ecfe1750a5..f4241137df 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -12,19 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( hflux_ffsl_hybrid_stencil_02, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - _shape, - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py index 85f9ca14c7..74c1bf9d7f 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py @@ -12,19 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( hflx_limiter_mo_stencil_01a, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - _shape, - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py index 8f7f4f11af..11417d34c4 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py @@ -12,8 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator import embedded as it_embedded from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( @@ -21,11 +19,8 @@ ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - _shape, as_1D_sparse_field, - constant_field, random_field, - zero_field, ) from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py index 95ced3a802..ed67d56456 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -13,14 +13,12 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( prep_gauss_quadrature_c_list_stencil, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, constant_field, random_field, zero_field, diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py index 77e574074a..315fbba080 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py @@ -12,17 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( prep_gauss_quadrature_c_stencil, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py index 8945975f34..dc2d8f1203 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py @@ -49,7 +49,6 @@ def recon_lsq_cell_c_svd_stencil_numpy( lsq_moments_9: np.ndarray, ) -> tuple[np.ndarray]: - p_cc_e = np.expand_dims(p_cc, axis=-1) lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) @@ -323,7 +322,7 @@ def test_recon_lsq_cell_c_svd_stencil(): ), }, ) - # co1 = np.asarray(p_coeff_1_dsl) + co1 = np.asarray(p_coeff_1_dsl) co2 = np.asarray(p_coeff_2_dsl) co3 = np.asarray(p_coeff_3_dsl) co4 = np.asarray(p_coeff_4_dsl) @@ -333,8 +332,7 @@ def test_recon_lsq_cell_c_svd_stencil(): co8 = np.asarray(p_coeff_8_dsl) co9 = np.asarray(p_coeff_9_dsl) co10 = np.asarray(p_coeff_10_dsl) - # assert np.allclose(ref_1, co1) - assert np.allclose(ref_10, co10) + assert np.allclose(ref_1, co1) assert np.allclose(ref_2, co2) assert np.allclose(ref_3, co3) assert np.allclose(ref_4, co4) @@ -343,3 +341,4 @@ def test_recon_lsq_cell_c_svd_stencil(): assert np.allclose(ref_7, co7) assert np.allclose(ref_8, co8) assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py index 6fff27fbce..82b5b57d3e 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py @@ -159,7 +159,6 @@ def test_upwind_hflux_miura3_stencil_01(): z_quad_vector_sum_10 = random_field(mesh, EdgeDim, KDim) p_mass_flx_e = random_field(mesh, EdgeDim, KDim) z_dreg_area = random_field(mesh, EdgeDim, KDim) - # cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) p_out_e_miura3 = zero_field(mesh, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py index ef9c229747..2c574b4855 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -19,7 +19,6 @@ ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( - _shape, random_field, random_mask, zero_field, diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py index b6d898e080..cae5ac6d30 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -19,12 +19,7 @@ upwind_hflux_miura_cycl_stencil_02, ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - _shape, - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -64,7 +59,6 @@ def test_upwind_hflux_miura_cycl_stencil_02(): nsub = int32(1) p_mass_flx_e = random_field(mesh, EdgeDim, KDim) geofac_div = random_field(mesh, CellDim, C2EDim) - geofac_div_field = as_1D_sparse_field(geofac_div, CEDim) z_rhofluxdiv_c = random_field(mesh, CellDim, KDim) z_tracer_mflx = random_field(mesh, EdgeDim, KDim) z_rho_now = random_field(mesh, CellDim, KDim) diff --git a/model/atmosphere/dycore/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py index 406c755810..7501aa0738 100644 --- a/model/atmosphere/dycore/tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/test_compute_airmass.py @@ -12,13 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np + from icon4py.model.atmosphere.dycore.compute_airmass import compute_airmass -from icon4py.model.common.dimension import EdgeDim, KDim, CellDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh - def compute_airmass_numpy( rho_in: np.array, ddqz_z_full_in: np.array, deepatmo_t1mc_in: np.array ) -> np.array: diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index d6f92a2cb8..74d7fbb7e4 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -350,27 +350,28 @@ class SimpleMeshData: c2e2c2e2c_table = np.asarray( [ - [15, 4, 3,12, 14, 1, 7, 6, 2], #1c - [16, 5, 4,12, 13, 2, 8, 7, 0], - [17, 3, 5,13, 14, 0, 6, 8, 1], + [15, 4, 3, 12, 14, 1, 7, 6, 2], # 1c + [16, 5, 4, 12, 13, 2, 8, 7, 0], + [17, 3, 5, 13, 14, 0, 6, 8, 1], [0, 6, 2, 17, 5, 9, 10, 15, 4], - [1, 7, 0, 15, 3, 16, 5, 10, 11], #5c + [1, 7, 0, 15, 3, 16, 5, 10, 11], # 5c [2, 8, 1, 4, 16, 17, 3, 9, 11], [3, 10, 9, 2, 0, 7, 13, 8, 12], [4, 11, 10, 0, 1, 8, 14, 6, 13], [5, 9, 11, 1, 2, 3, 12, 7, 14], - [6, 12, 8, 5, 11, 3, 10, 16, 15], #10c + [6, 12, 8, 5, 11, 3, 10, 16, 15], # 10c [7, 13, 6, 3, 9, 4, 11, 16, 17], [8, 14, 7, 4, 10, 5, 9, 15, 17], [9, 16, 15, 8, 6, 1, 13, 0, 14], [10, 17, 16, 6, 7, 2, 14, 1, 12], - [11, 15, 17, 7, 8, 2, 13, 0, 12], #15c - [12, 0, 14,11, 17, 9, 16, 3, 4], + [11, 15, 17, 7, 8, 2, 13, 0, 12], # 15c + [12, 0, 14, 11, 17, 9, 16, 3, 4], [13, 1, 12, 9, 15, 10, 17, 4, 5], - [14, 2, 13,10, 16, 5, 3, 11, 15], + [14, 2, 13, 10, 16, 5, 3, 11, 15], ] ) + class SimpleMesh: _DEFAULT_K_LEVEL = 10 @@ -413,7 +414,6 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): E2C2EODim: self.n_e2c2eO, E2C2EDim: self.n_e2c2e, V2CDim: self.n_v2c, - V2CDim: self.n_v2c, KDim: self.k_level, VertexDim: self.n_vertices, V2EDim: self.n_v2e, From d45cdc25f99bab8f57e05d162dcf68bcb71c46ff Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 28 Aug 2023 09:02:40 +0200 Subject: [PATCH 091/105] Pre-commit run --- .../tests/test_btraj_dreg_stencil_01.py | 4 ++-- ...t_mo_advection_traj_btraj_compute_o1_dsl.py | 6 ++++-- ...late_nabla2_and_smag_coefficients_for_vn.py | 4 +++- .../tests/test_calculate_nabla2_of_theta.py | 8 ++++++-- .../tests/test_mo_solve_nonhydro_stencil_05.py | 8 ++++++-- .../tests/test_mo_solve_nonhydro_stencil_10.py | 8 ++++++-- ..._nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 ++++-- .../tests/test_mo_solve_nonhydro_stencil_39.py | 14 +++++++++++--- .../tests/test_mo_solve_nonhydro_stencil_40.py | 18 ++++++++++++++---- .../test_mo_velocity_advection_stencil_02.py | 12 +++++++++--- .../test_mo_velocity_advection_stencil_10.py | 8 ++++++-- 11 files changed, 71 insertions(+), 25 deletions(-) diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py index b136704f95..aeba344431 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py @@ -17,8 +17,8 @@ btraj_dreg_stencil_01, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import random_field -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh, zero_field +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def btraj_dreg_stencil_01_numpy( diff --git a/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 6f7cd10dfc..4522ed5d9c 100644 --- a/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -71,10 +71,12 @@ def reference( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index a9a31ad410..89a1166fed 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -117,7 +117,9 @@ def reference( + v_vert_e2c2v[:, 3] * primal_normal_vert_y[:, 3] ) - kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - (dvt_tang * inv_primal_edge_length) + kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - ( + dvt_tang * inv_primal_edge_length + ) kh_smag_2 = kh_smag_2 * kh_smag_2 diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py index 2fceed17ed..00cc2ee9d7 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py @@ -31,10 +31,14 @@ class TestCalculateNabla2OfTheta(StencilTest): OUTPUTS = ("z_temp",) @staticmethod - def reference(mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs) -> np.array: + def reference( + mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs + ) -> np.array: geofac_div = geofac_div.reshape(mesh.c2e.shape) geofac_div = np.expand_dims(geofac_div, axis=-1) - z_temp = np.sum(z_nabla2_e[mesh.c2e] * geofac_div, axis=1) # sum along edge dimension + z_temp = np.sum( + z_nabla2_e[mesh.c2e] * geofac_div, axis=1 + ) # sum along edge dimension return dict(z_temp=z_temp) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py index 8eab825615..335f8ff3c1 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py @@ -30,9 +30,13 @@ class TestMoSolveNonhydroStencil05(StencilTest): OUTPUTS = ("z_exner_ic",) @staticmethod - def reference(mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs) -> np.array: + def reference( + mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs + ) -> np.array: z_exner_ex_pr_offset_1 = np.roll(z_exner_ex_pr, shift=1, axis=1) - z_exner_ic = wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 + z_exner_ic = ( + wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 + ) return dict(z_exner_ic=z_exner_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py index d66d399fa2..c2d1d80e2b 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py @@ -55,7 +55,9 @@ def reference( z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now_offset + wgt_nnew_rth * rho_var_offset - z_theta_tavg_m1 = wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset + z_theta_tavg_m1 = ( + wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset + ) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( @@ -65,7 +67,9 @@ def reference( ) z_theta_v_pr_mc_m1 = z_theta_tavg_m1 - theta_ref_mc_offset z_theta_v_pr_mc = z_theta_tavg - theta_ref_mc - z_theta_v_pr_ic = wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 + z_theta_v_pr_ic = ( + wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 + ) theta_v_ic = ( wgtfac_c * z_theta_tavg + (1 - wgtfac_c) * z_theta_tavg_m1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 686f58cf7f..d78bf4e4cf 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -52,10 +52,12 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index af619ab8cf..8073cbdf72 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -31,13 +31,21 @@ class TestMoSolveNonhydroStencil39(StencilTest): @staticmethod def reference( - mesh, e_bln_c_s: np.array, z_w_concorr_me: np.array, wgtfac_c: np.array, **kwargs + mesh, + e_bln_c_s: np.array, + z_w_concorr_me: np.array, + wgtfac_c: np.array, + **kwargs, ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_mc_m0 = np.sum(e_bln_c_s * z_w_concorr_me[mesh.c2e], axis=1) - z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) - w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 + z_w_concorr_mc_m1 = np.sum( + e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 + ) + w_concorr_c = ( + wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 + ) return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index fb9f0ca91f..b89e93ba3a 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -31,15 +31,25 @@ class TestMoSolveNonhydroStencil40(StencilTest): @staticmethod def reference( - mesh, e_bln_c_s: np.array, z_w_concorr_me: np.array, wgtfacq_c: np.array, **kwargs + mesh, + e_bln_c_s: np.array, + z_w_concorr_me: np.array, + wgtfacq_c: np.array, + **kwargs, ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_me_offset_2 = np.roll(z_w_concorr_me, shift=2, axis=1) z_w_concorr_me_offset_3 = np.roll(z_w_concorr_me, shift=3, axis=1) - z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) - z_w_concorr_mc_m2 = np.sum(e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1) - z_w_concorr_mc_m3 = np.sum(e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1) + z_w_concorr_mc_m1 = np.sum( + e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 + ) + z_w_concorr_mc_m2 = np.sum( + e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1 + ) + z_w_concorr_mc_m3 = np.sum( + e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1 + ) w_concorr_c = ( np.roll(wgtfacq_c, shift=1, axis=1) * z_w_concorr_mc_m1 + np.roll(wgtfacq_c, shift=2, axis=1) * z_w_concorr_mc_m2 diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index b677e09773..c8b57645cc 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -30,18 +30,24 @@ class TestMoVelocityAdvectionStencil02VnIe(StencilTest): OUTPUTS = ("vn_ie", "z_kin_hor_e") @staticmethod - def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: + def mo_velocity_advection_stencil_02_vn_ie_numpy( + wgtfac_e: np.array, vn: np.array + ) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie @staticmethod - def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: + def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( + vn: np.array, vt: np.array + ) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e @classmethod - def reference(cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs) -> dict: + def reference( + cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs + ) -> dict: vn_ie = cls.mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e, vn) z_kin_hor_e = cls.mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn, vt) return dict(vn_ie=vn_ie, z_kin_hor_e=z_kin_hor_e) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index d69e70a4f2..997d4d450e 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -30,9 +30,13 @@ class TestMoVelocityAdvectionStencil10(StencilTest): OUTPUTS = ("w_concorr_c",) @staticmethod - def reference(mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs) -> np.array: + def reference( + mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs + ) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + w_concorr_c = ( + wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + ) return dict(w_concorr_c=w_concorr_c) @pytest.fixture From f42ae385ac001cb8179a80eece8db7017208d5b6 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 8 Sep 2023 09:03:29 +0200 Subject: [PATCH 092/105] Adapt pre-commit and QA fixes for advection component --- .github/workflows/icon4py-qa.yml | 24 ++++ model/atmosphere/advection/.flake8 | 42 +++++++ .../advection/.pre-commit-config.yaml | 114 ++++++++++++++++++ model/atmosphere/advection/pyproject.toml | 75 ++++++------ 4 files changed, 218 insertions(+), 37 deletions(-) create mode 100644 model/atmosphere/advection/.flake8 create mode 100644 model/atmosphere/advection/.pre-commit-config.yaml diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 18be464b94..0b7bfdb7eb 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -69,3 +69,27 @@ jobs: - name: Run checks icon4py-model-atmosphere-dycore run: | pre-commit run --config model/atmosphere/dycore/.pre-commit-config.yaml --all-files + + pre-commit-icon4py-model-atmosphere-advection: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + cache: 'pip' + cache-dependency-path: | + **/pyproject.toml + **/base-requirements.txt + **/base-requirements-dev.txt + **/requirements.txt + **/requirements-dev.txt + - name: Install icon4py-model-atmosphere-advection + working-directory: model/atmosphere/advection + run: | + python -m pip install --upgrade pip setuptools wheel + python -m pip install -r ./requirements-dev.txt + - name: Run checks icon4py-model-atmosphere-advection + run: | + pre-commit run --config model/atmosphere/advection/.pre-commit-config.yaml --all-files diff --git a/model/atmosphere/advection/.flake8 b/model/atmosphere/advection/.flake8 new file mode 100644 index 0000000000..31cecff5ab --- /dev/null +++ b/model/atmosphere/advection/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/atmosphere/advection/.pre-commit-config.yaml b/model/atmosphere/advection/.pre-commit-config.yaml new file mode 100644 index 0000000000..a5f7495843 --- /dev/null +++ b/model/atmosphere/advection/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/atmosphere/advection/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/atmosphere/advection/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/atmosphere/advection/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/atmosphere/advection/.flake8, model/atmosphere/advection/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index 268c47eed0..65d527bff9 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -1,36 +1,36 @@ [build-system] -requires = ["setuptools>=61.0", "wheel>=0.40.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] [project] -name = "icon4py-atmosphere-advection" -description = "ICON advection" -readme = "README.md" -requires-python = ">=3.10" -license = {file = "LICENSE"} authors = [ - {email = "gridtools@cscs.ch"}, - {name = "ETH Zurich"} + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} ] classifiers = [ - "Development Status :: 3 - Alpha", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", - "Operating System :: POSIX", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: Implementation :: CPython", - "Topic :: Scientific/Engineering :: Atmospheric Science", - "Topic :: Scientific/Engineering :: Mathematics", - "Topic :: Scientific/Engineering :: Physics" + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1", - "icon4py_common>=0.0.5", + "gt4py>=1.0.1", + "icon4py_common>=0.0.5" ] +description = "ICON advection." dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-atmosphere-advection" +readme = "README.md" +requires-python = ">=3.10" [project.urls] repository = "https://github.com/C2SM/icon4py" @@ -71,40 +71,41 @@ exclude_lines = [ ignore_errors = true [tool.coverage.run] -parallel = true branch = true +parallel = true source_pkgs = ['advection'] [tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py'] +known_third_party = ['gt4py'] lexicographical = true line_length = 100 # It should be the same as in `tool.black.line-length` above lines_after_imports = 2 +# known_tests = ['tests_utils'] +multi_line_output = 3 profile = 'black' -sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'LOCALFOLDER'] skip_gitignore = true skip_glob = ['*.venv/**', '_local/**'] -known_first_party = ['icon4py.model'] -known_third_party = ['gt4py'] -multi_line_output = 3 use_parentheses = true -include_trailing_comma = true -force_grid_wrap = 0 [tool.mypy] -install_types = true -non_interactive = true -exclude = [ - '^tests/*.py', -] disallow_incomplete_defs = true disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] ignore_missing_imports = false implicit_reexport = true -warn_unused_configs = true -warn_unused_ignores = true -warn_redundant_casts = true +install_types = true +non_interactive = true show_column_numbers = true show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true [tool.pytest] From a76337f6ceb493ba3f2278176ddb3d19f02091be Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 8 Sep 2023 16:19:47 +0200 Subject: [PATCH 093/105] Pre-commit run with new configs --- .../advection/btraj_dreg_stencil_03.py | 56 ++------ .../divide_flux_area_list_stencil_01.py | 128 +++++------------ .../divide_flux_area_list_stencil_02.py | 32 ++--- .../advection/face_val_ppm_stencil_01.py | 3 +- .../advection/hflx_limiter_mo_stencil_01b.py | 26 +--- .../advection/hflx_limiter_mo_stencil_03.py | 8 +- .../advection/hflx_limiter_pd_stencil_01.py | 13 +- .../advection/hor_adv_stencil_01.py | 10 +- .../mo_advection_traj_btraj_compute_o1_dsl.py | 6 +- .../prep_gauss_quadrature_c_list_stencil.py | 4 +- .../prep_gauss_quadrature_c_stencil.py | 8 +- .../advection/recon_lsq_cell_c_stencil.py | 8 +- .../advection/recon_lsq_cell_c_svd_stencil.py | 8 +- .../advection/step_advection_stencil_03.py | 4 +- .../advection/step_advection_stencil_04.py | 4 +- .../upwind_hflux_miura_cycl_stencil_02.py | 4 +- .../upwind_hflux_miura_cycl_stencil_03a.py | 4 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 4 +- .../upwind_hflux_miura_stencil_01.py | 12 +- .../advection/vert_adv_stencil_01.py | 3 +- .../tests/test_btraj_dreg_stencil_01.py | 4 +- .../tests/test_btraj_dreg_stencil_02.py | 10 +- .../tests/test_btraj_dreg_stencil_03.py | 66 +++------ .../test_divide_flux_area_list_stencil_01.py | 132 +++++------------- .../test_divide_flux_area_list_stencil_02.py | 6 +- .../tests/test_face_val_ppm_stencil_01.py | 30 +--- .../tests/test_face_val_ppm_stencil_02.py | 13 +- .../tests/test_face_val_ppm_stencil_02a.py | 13 +- .../tests/test_face_val_ppm_stencil_02b.py | 4 +- .../tests/test_face_val_ppm_stencil_02c.py | 4 +- .../tests/test_face_val_ppm_stencil_05.py | 4 +- .../test_hflux_ffsl_hybrid_stencil_01a.py | 6 +- .../tests/test_hflx_limiter_mo_stencil_01b.py | 5 +- .../tests/test_hflx_limiter_mo_stencil_02.py | 18 +-- .../tests/test_hflx_limiter_mo_stencil_04.py | 4 +- .../tests/test_hflx_limiter_pd_stencil_01.py | 10 +- .../tests/test_hflx_limiter_pd_stencil_02.py | 4 +- .../tests/test_hor_adv_stencil_01.py | 9 +- ..._mo_advection_traj_btraj_compute_o1_dsl.py | 6 +- ...st_prep_gauss_quadrature_c_list_stencil.py | 10 +- .../test_prep_gauss_quadrature_c_stencil.py | 4 +- .../tests/test_rbf_intp_edge_stencil_01.py | 4 +- .../tests/test_recon_lsq_cell_c_stencil.py | 14 +- .../test_recon_lsq_cell_c_svd_stencil.py | 10 +- .../test_recon_lsq_cell_l_svd_stencil.py | 6 +- .../tests/test_step_advection_stencil_01.py | 7 +- .../tests/test_step_advection_stencil_02.py | 9 +- .../tests/test_step_advection_stencil_03.py | 4 +- .../tests/test_step_advection_stencil_04.py | 4 +- .../test_upwind_hflux_miura3_stencil_01.py | 6 +- ...test_upwind_hflux_miura_cycl_stencil_01.py | 6 +- ...test_upwind_hflux_miura_cycl_stencil_02.py | 8 +- ...est_upwind_hflux_miura_cycl_stencil_03b.py | 4 +- .../test_upwind_hflux_miura_stencil_01.py | 6 +- .../tests/test_upwind_vflux_ppm_stencil_01.py | 4 +- .../tests/test_vert_adv_stencil_01.py | 9 +- .../tests/test_vlimit_prbl_sm_stencil_01.py | 4 +- .../tests/test_vlimit_prbl_sm_stencil_02.py | 10 +- .../dycore/apply_diffusion_to_vn.py | 4 +- ...ute_horizontal_gradients_for_turbulance.py | 8 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 4 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 4 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 4 +- .../calculate_diagnostics_for_turbulence.py | 4 +- ..._coefficients_for_grid_point_cold_pools.py | 4 +- ...ate_horizontal_gradients_for_turbulence.py | 4 +- .../dycore/calculate_nabla2_for_theta.py | 4 +- ...n_coefficient_for_grid_point_cold_pools.py | 4 +- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 4 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_02.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_10.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 +- .../dycore/mo_solve_nonhydro_stencil_17.py | 5 +- .../dycore/mo_solve_nonhydro_stencil_18.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_20.py | 10 +- .../dycore/mo_solve_nonhydro_stencil_22.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_24.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_39.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_40.py | 12 +- .../dycore/mo_solve_nonhydro_stencil_42.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_55.py | 13 +- .../dycore/mo_solve_nonhydro_stencil_56_63.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_58.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_60.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_65.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_67.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_68.py | 5 +- .../mo_velocity_advection_stencil_04.py | 4 +- .../mo_velocity_advection_stencil_07.py | 4 +- .../mo_velocity_advection_stencil_10.py | 4 +- .../mo_velocity_advection_stencil_14.py | 8 +- .../mo_velocity_advection_stencil_16.py | 8 +- .../mo_velocity_advection_stencil_19.py | 5 +- ...fusion_nabla_of_theta_over_steep_points.py | 15 +- .../dycore/update_theta_and_exner.py | 4 +- ...st_apply_nabla2_and_nabla4_global_to_vn.py | 8 +- .../dycore/tests/test_apply_nabla2_to_w.py | 4 +- ...st_calculate_diagnostics_for_turbulence.py | 8 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 4 +- .../tests/test_calculate_nabla2_of_theta.py | 8 +- .../dycore/tests/test_compute_airmass.py | 4 +- .../test_mo_solve_nonhydro_stencil_02.py | 4 +- .../test_mo_solve_nonhydro_stencil_04.py | 9 +- .../test_mo_solve_nonhydro_stencil_05.py | 8 +- .../test_mo_solve_nonhydro_stencil_06.py | 4 +- .../test_mo_solve_nonhydro_stencil_10.py | 8 +- .../test_mo_solve_nonhydro_stencil_15.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 +- .../test_mo_solve_nonhydro_stencil_19.py | 5 +- .../test_mo_solve_nonhydro_stencil_20.py | 8 +- .../test_mo_solve_nonhydro_stencil_21.py | 4 +- .../test_mo_solve_nonhydro_stencil_25.py | 4 +- .../test_mo_solve_nonhydro_stencil_26.py | 4 +- .../test_mo_solve_nonhydro_stencil_28.py | 4 +- .../test_mo_solve_nonhydro_stencil_29.py | 4 +- .../test_mo_solve_nonhydro_stencil_39.py | 8 +- .../test_mo_solve_nonhydro_stencil_40.py | 12 +- .../test_mo_solve_nonhydro_stencil_42.py | 4 +- .../test_mo_solve_nonhydro_stencil_44.py | 4 +- .../test_mo_solve_nonhydro_stencil_47.py | 4 +- .../test_mo_solve_nonhydro_stencil_51.py | 4 +- .../test_mo_solve_nonhydro_stencil_52.py | 8 +- .../test_mo_solve_nonhydro_stencil_54.py | 4 +- .../test_mo_solve_nonhydro_stencil_55.py | 5 +- .../test_mo_solve_nonhydro_stencil_58.py | 4 +- .../test_mo_solve_nonhydro_stencil_60.py | 4 +- .../test_mo_solve_nonhydro_stencil_62.py | 4 +- .../test_mo_solve_nonhydro_stencil_65.py | 4 +- .../test_mo_solve_nonhydro_stencil_68.py | 5 +- .../test_mo_velocity_advection_stencil_02.py | 12 +- .../test_mo_velocity_advection_stencil_08.py | 4 +- .../test_mo_velocity_advection_stencil_09.py | 4 +- .../test_mo_velocity_advection_stencil_10.py | 8 +- .../test_mo_velocity_advection_stencil_13.py | 4 +- .../test_mo_velocity_advection_stencil_19.py | 5 +- ...fusion_nabla_of_theta_over_steep_points.py | 7 +- 137 files changed, 297 insertions(+), 958 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 8235d1b611..6dce27f72a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -66,54 +66,26 @@ def _btraj_dreg_stencil_03( pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y - pos_dreg_vert_c_2_x = ( - where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_2_y = ( - where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y - ) + pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + pos_dreg_vert_c_2_y = where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y - pos_dreg_vert_c_4_x = ( - where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_4_y = ( - where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - ) + pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - pn_cell_1 = where( - lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1]) - ) - pn_cell_2 = where( - lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1]) - ) + pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) + pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) - p_coords_dreg_v_1_lon_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 - ) - p_coords_dreg_v_2_lon_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 - ) - p_coords_dreg_v_3_lon_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 - ) - p_coords_dreg_v_4_lon_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - ) - p_coords_dreg_v_1_lat_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 - ) - p_coords_dreg_v_2_lat_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 - ) - p_coords_dreg_v_3_lat_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 - ) - p_coords_dreg_v_4_lat_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 - ) + p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + p_coords_dreg_v_4_lat_dsl = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 return ( p_cell_idx, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index f240bc0ccf..a753bd2380 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -112,9 +112,7 @@ def line_intersect( m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) - intersect_1 = ( - line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon - ) / (m1 - m2) + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) return intersect_1, intersect_2 @@ -258,12 +256,8 @@ def _divide_flux_area_list_stencil_01( ) # Case 1 - patch 0 - dreg_patch0_1_lon_dsl = where( - mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = where( - mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = where( mask_case1, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), @@ -324,12 +318,8 @@ def _divide_flux_area_list_stencil_01( # ------------------------------------------------- Case 2a mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool # Case 2a - patch 0 - dreg_patch0_1_lon_dsl = where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = where( mask_case2a, where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), @@ -340,12 +330,8 @@ def _divide_flux_area_list_stencil_01( where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = where( - mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl - ) - dreg_patch0_3_lat_dsl = where( - mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl - ) + dreg_patch0_3_lon_dsl = where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = where( mask_case2a, where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), @@ -357,18 +343,10 @@ def _divide_flux_area_list_stencil_01( dreg_patch0_4_lat_dsl, ) # Case 2a - patch 1 - dreg_patch1_1_lon_vmask = where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) - dreg_patch1_4_lon_vmask = where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask - ) - dreg_patch1_4_lat_vmask = where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask - ) + dreg_patch1_1_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) dreg_patch1_2_lon_vmask = where( mask_case2a, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), @@ -402,12 +380,8 @@ def _divide_flux_area_list_stencil_01( # -------------------------------------------------- Case 2b mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool # Case 2b - patch 0 - dreg_patch0_1_lon_dsl = where( - mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = where( - mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = where( mask_case2b, where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -440,18 +414,10 @@ def _divide_flux_area_list_stencil_01( dreg_patch1_4_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) # Case 2b - patch 2 - dreg_patch2_1_lon_vmask = where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) - dreg_patch2_4_lon_vmask = where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask - ) - dreg_patch2_4_lat_vmask = where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask - ) + dreg_patch2_1_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) dreg_patch2_2_lon_vmask = where( mask_case2b, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), @@ -507,12 +473,8 @@ def _divide_flux_area_list_stencil_01( tri_line1_p2_lat, ) # Case 3a - patch 0 - dreg_patch0_1_lon_dsl = where( - mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = where( mask_case3a, where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -536,12 +498,8 @@ def _divide_flux_area_list_stencil_01( dreg_patch0_4_lat_dsl, ) # Case 3a - patch 1 - dreg_patch1_1_lon_vmask = where( - mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) + dreg_patch1_1_lon_vmask = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) dreg_patch1_2_lon_vmask = where( mask_case3a, where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), @@ -552,12 +510,8 @@ def _divide_flux_area_list_stencil_01( where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = where( - mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask - ) - dreg_patch1_3_lat_vmask = where( - mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask - ) + dreg_patch1_3_lon_vmask = where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) dreg_patch1_4_lon_vmask = where( mask_case3a, where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), @@ -602,18 +556,10 @@ def _divide_flux_area_list_stencil_01( tri_line2_p2_lat, ) # Case 3b - patch 0 - dreg_patch0_1_lon_dsl = where( - mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = where( - mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) - dreg_patch0_4_lon_dsl = where( - mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl - ) - dreg_patch0_4_lat_dsl = where( - mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl - ) + dreg_patch0_1_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) dreg_patch0_2_lon_dsl = where( mask_case3b, where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), @@ -644,12 +590,8 @@ def _divide_flux_area_list_stencil_01( dreg_patch1_4_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 - dreg_patch2_1_lon_vmask = where( - mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = where( - mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) + dreg_patch2_1_lon_vmask = where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) dreg_patch2_2_lon_vmask = where( mask_case3b, where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), @@ -660,12 +602,8 @@ def _divide_flux_area_list_stencil_01( where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = where( - mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask - ) - dreg_patch2_3_lat_vmask = where( - mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask - ) + dreg_patch2_3_lon_vmask = where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) dreg_patch2_4_lon_vmask = where( mask_case3b, where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), @@ -680,9 +618,7 @@ def _divide_flux_area_list_stencil_01( # --------------------------------------------- Case 4 # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions # can't be overwritten by this new condition. - indices_previously_matched = ( - mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1 - ) + indices_previously_matched = mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1 # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices mask_case4 = famask_bool & (not indices_previously_matched) # Case 4 - patch 0 - no change diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py index 4cd783407b..08650dade8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py @@ -121,30 +121,14 @@ def _divide_flux_area_list_stencil_02( # Store global index of the underlying grid cell # Adapt dimensions to fit ofr multiple levels - butterfly_idx_patch1_vnpos_3d = broadcast( - butterfly_idx_patch1_vnpos, (EdgeDim, KDim) - ) - butterfly_idx_patch1_vnneg_3d = broadcast( - butterfly_idx_patch1_vnneg, (EdgeDim, KDim) - ) - butterfly_idx_patch2_vnpos_3d = broadcast( - butterfly_idx_patch2_vnpos, (EdgeDim, KDim) - ) - butterfly_idx_patch2_vnneg_3d = broadcast( - butterfly_idx_patch2_vnneg, (EdgeDim, KDim) - ) - butterfly_blk_patch1_vnpos_3d = broadcast( - butterfly_blk_patch1_vnpos, (EdgeDim, KDim) - ) - butterfly_blk_patch1_vnneg_3d = broadcast( - butterfly_blk_patch1_vnneg, (EdgeDim, KDim) - ) - butterfly_blk_patch2_vnpos_3d = broadcast( - butterfly_blk_patch2_vnpos, (EdgeDim, KDim) - ) - butterfly_blk_patch2_vnneg_3d = broadcast( - butterfly_blk_patch2_vnneg, (EdgeDim, KDim) - ) + butterfly_idx_patch1_vnpos_3d = broadcast(butterfly_idx_patch1_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch1_vnneg_3d = broadcast(butterfly_idx_patch1_vnneg, (EdgeDim, KDim)) + butterfly_idx_patch2_vnpos_3d = broadcast(butterfly_idx_patch2_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch2_vnneg_3d = broadcast(butterfly_idx_patch2_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch1_vnpos_3d = broadcast(butterfly_blk_patch1_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch1_vnneg_3d = broadcast(butterfly_blk_patch1_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch2_vnpos_3d = broadcast(butterfly_blk_patch2_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch2_vnneg_3d = broadcast(butterfly_blk_patch2_vnneg, (EdgeDim, KDim)) patch1_cell_idx_vmask = where( famask_bool, where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py index 5ffaabbe96..de605d3dd4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -45,8 +45,7 @@ def _face_val_ppm_stencil_01b( zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) z_slope = ( - p_cellhgt_mc_now - / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) ) * ( (2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + (p_cellhgt_mc_now + 2.0 * p_cellhgt_mc_now) * zfac_m1 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py index fcc5e26405..0c9c943249 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py @@ -12,23 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import ( - Field, - broadcast, - maximum, - minimum, - neighbor_sum, -) +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, neighbor_sum -from icon4py.model.common.dimension import ( - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator @@ -54,15 +40,11 @@ def _hflx_limiter_mo_stencil_01b( z_mflx_anti_3 = p_dtime * geofac_div(C2CE[2]) / p_rhodz_new * z_anti(C2E[2]) z_mflx_anti_in = -1.0 * ( - minimum(zero, z_mflx_anti_1) - + minimum(zero, z_mflx_anti_2) - + minimum(zero, z_mflx_anti_3) + minimum(zero, z_mflx_anti_1) + minimum(zero, z_mflx_anti_2) + minimum(zero, z_mflx_anti_3) ) z_mflx_anti_out = ( - maximum(zero, z_mflx_anti_1) - + maximum(zero, z_mflx_anti_2) - + maximum(zero, z_mflx_anti_3) + maximum(zero, z_mflx_anti_1) + maximum(zero, z_mflx_anti_2) + maximum(zero, z_mflx_anti_3) ) z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div(C2CE), axis=C2EDim) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py index 2dea287c56..a7184d1055 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py @@ -25,12 +25,8 @@ def _hflx_limiter_mo_stencil_03_min_max( beta_fct: float, r_beta_fct: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_max = beta_fct * maximum( - max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max - ) - z_min = r_beta_fct * minimum( - min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min - ) + z_max = beta_fct * maximum(max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max) + z_min = r_beta_fct * minimum(min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min) return z_max, z_min diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py index cc51b56438..39c2534870 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py @@ -14,14 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum -from icon4py.model.common.dimension import ( - C2CE, - C2E, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.model.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim @field_operator @@ -39,9 +32,7 @@ def _hflx_limiter_pd_stencil_01( pm_1 = maximum(zero, p_mflx_tracer_h(C2E[1]) * geofac_div(C2CE[1]) * p_dtime) pm_2 = maximum(zero, p_mflx_tracer_h(C2E[2]) * geofac_div(C2CE[2]) * p_dtime) p_m = pm_0 + pm_1 + pm_2 - r_m = minimum( - broadcast(1.0, (CellDim, KDim)), (p_cc * p_rhodz_now) / (p_m + dbl_eps) - ) + r_m = minimum(broadcast(1.0, (CellDim, KDim)), (p_cc * p_rhodz_now) / (p_m + dbl_eps)) return r_m diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py index b6ca46cbe8..59ad061a9d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py @@ -14,15 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import ( - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py index fac30e51fd..572b0f8d64 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py @@ -45,13 +45,11 @@ def _mo_advection_traj_btraj_compute_o1_dsl( p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py index 7fec10aad9..957ee74a09 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py @@ -251,9 +251,7 @@ def _prep_gauss_quadrature_c_list_stencil( + shape_func_4_4 * p_coords_dreg_v_4_y ) - p_quad_vector_sum_1 = ( - wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - ) + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 p_quad_vector_sum_2 = ( wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py index 228ed10d93..d75ed51a38 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py @@ -215,9 +215,7 @@ def _prep_gauss_quadrature_c_stencil( + shape_func_4_4 * p_coords_dreg_v_4_y ) - p_quad_vector_sum_1 = ( - wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - ) + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 p_quad_vector_sum_2 = ( wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x @@ -274,9 +272,7 @@ def _prep_gauss_quadrature_c_stencil( ) z_area = p_quad_vector_sum_1 - p_dreg_area_out = where( - z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area)) - ) + p_dreg_area_out = where(z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area))) return ( p_quad_vector_sum_1, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py index abea9c0f40..4866b8f2dd 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py @@ -16,13 +16,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.model.common.dimension import ( - C2CECEC, - C2E2C2E2C, - CECECDim, - CellDim, - KDim, -) +from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim sys.setrecursionlimit(6000) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py index edffe26fcb..b16a6d96ab 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -14,13 +14,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program -from icon4py.model.common.dimension import ( - C2CECEC, - C2E2C2E2C, - CECECDim, - CellDim, - KDim, -) +from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim @field_operator diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py index 85b8ac433a..b1e775c62d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py @@ -34,6 +34,4 @@ def step_advection_stencil_03( p_tracer_new: Field[[CellDim, KDim], float], p_dtime: float, ): - _step_advection_stencil_03( - p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new - ) + _step_advection_stencil_03(p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py index a274a355b2..52e0c79402 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py @@ -34,6 +34,4 @@ def step_advection_stencil_04( opt_ddt_tracer_adv: Field[[CellDim, KDim], float], p_dtime: float, ): - _step_advection_stencil_04( - p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv - ) + _step_advection_stencil_04(p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index c4ae3fa8ec..b0026495f8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -44,9 +44,7 @@ def _upwind_hflux_miura_cycl_stencil_02( z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out - z_tracer_new_dsl = ( - z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl - ) / z_rho_new_dsl + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py index a8bcf8117f..53cdc339d1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -33,6 +33,4 @@ def upwind_hflux_miura_cycl_stencil_03a( z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], p_out_e: Field[[EdgeDim, KDim], float], ): - _upwind_hflux_miura_cycl_stencil_03a( - z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e) - ) + _upwind_hflux_miura_cycl_stencil_03a(z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py index 4f5f73ac91..4b5c73e401 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -23,9 +23,7 @@ def _upwind_hflux_miura_cycl_stencil_03b( z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float( - 3 - ) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) return p_out_e diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py index aa972a9285..3637ba94ad 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -30,17 +30,11 @@ def _upwind_hflux_miura_stencil_01( ) -> Field[[EdgeDim, KDim], float]: p_out_e = ( - where( - cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0]) - ) + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) + distv_bary_1 - * where( - cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0]) - ) + * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) + distv_bary_2 - * where( - cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0]) - ) + * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) ) * p_mass_flx_e return p_out_e diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 80d7f9d0e2..78736914db 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -29,8 +29,7 @@ def _vert_adv_stencil_01( ) -> Field[[CellDim, KDim], float]: tracer_new = ( tracer_now * rhodz_now - + p_dtime - * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) ) / rhodz_new return tracer_new diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py index aeba344431..8d162d7b9e 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import ( - btraj_dreg_stencil_01, -) +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py index ecd96ed72a..9e59a148c5 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py @@ -15,15 +15,9 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import ( - btraj_dreg_stencil_02, -) +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 from icon4py.model.common.dimension import ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py index 4c3a0e088a..712d4ecbbe 100644 --- a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py @@ -15,15 +15,9 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import ( - btraj_dreg_stencil_03, -) +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - constant_field, - random_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, constant_field, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -77,54 +71,26 @@ def btraj_dreg_stencil_03_numpy( pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y - pos_dreg_vert_c_2_x = ( - np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_2_y = ( - np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y - ) + pos_dreg_vert_c_2_x = np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + pos_dreg_vert_c_2_y = np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y - pos_dreg_vert_c_4_x = ( - np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_4_y = ( - np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - ) + pos_dreg_vert_c_4_x = np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + pos_dreg_vert_c_4_y = np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - pn_cell_1 = np.where( - lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1] - ) - pn_cell_2 = np.where( - lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1] - ) + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) - p_coords_dreg_v_1_lon_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 - ) - p_coords_dreg_v_2_lon_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 - ) - p_coords_dreg_v_3_lon_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 - ) - p_coords_dreg_v_4_lon_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - ) - p_coords_dreg_v_1_lat_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 - ) - p_coords_dreg_v_2_lat_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 - ) - p_coords_dreg_v_3_lat_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 - ) - p_coords_dreg_v_4_lat_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 - ) + p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + p_coords_dreg_v_4_lat_dsl = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 return ( p_cell_idx, diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py index 0402167fa0..82cba3671a 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py @@ -115,9 +115,7 @@ def line_intersect( m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) - intersect_1 = ( - line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon - ) / (m1 - m2) + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) return intersect_1, intersect_2 @@ -220,9 +218,7 @@ def divide_flux_area_list_stencil_01_numpy( ) famask_bool = np.where(famask_int == int32(1), True, False) # ------------------------------------------------- Case 1 - mask_case1 = np.logical_and.reduce( - [lintersect_line1, lintersect_line2, famask_bool] - ) + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) ps1_x, ps1_y = line_intersect( fl_line_p1_lon, fl_line_p1_lat, @@ -245,12 +241,8 @@ def divide_flux_area_list_stencil_01_numpy( ) # Case 1 - patch 0 - dreg_patch0_1_lon_dsl = np.where( - mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = np.where( - mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = np.where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = np.where( mask_case1, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), @@ -313,12 +305,8 @@ def divide_flux_area_list_stencil_01_numpy( [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] ) # Case 2a - patch 0 - dreg_patch0_1_lon_dsl = np.where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = np.where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = np.where( mask_case2a, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), @@ -329,12 +317,8 @@ def divide_flux_area_list_stencil_01_numpy( np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = np.where( - mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl - ) - dreg_patch0_3_lat_dsl = np.where( - mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl - ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) dreg_patch0_4_lon_dsl = np.where( mask_case2a, np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), @@ -346,18 +330,10 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch0_4_lat_dsl, ) # Case 2a - patch 1 - dreg_patch1_1_lon_vmask = np.where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = np.where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) - dreg_patch1_4_lon_vmask = np.where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask - ) - dreg_patch1_4_lat_vmask = np.where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask - ) + dreg_patch1_1_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) dreg_patch1_2_lon_vmask = np.where( mask_case2a, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), @@ -393,12 +369,8 @@ def divide_flux_area_list_stencil_01_numpy( [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] ) # Case 2b - patch 0 - dreg_patch0_1_lon_dsl = np.where( - mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = np.where( - mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = np.where( mask_case2b, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -431,18 +403,10 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch1_4_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) # Case 2b - patch 2 - dreg_patch2_1_lon_vmask = np.where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = np.where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) - dreg_patch2_4_lon_vmask = np.where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask - ) - dreg_patch2_4_lat_vmask = np.where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask - ) + dreg_patch2_1_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) dreg_patch2_2_lon_vmask = np.where( mask_case2b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), @@ -498,12 +462,8 @@ def divide_flux_area_list_stencil_01_numpy( tri_line1_p2_lat, ) # Case 3a - patch 0 - dreg_patch0_1_lon_dsl = np.where( - mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = np.where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) dreg_patch0_2_lon_dsl = np.where( mask_case3a, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), @@ -527,12 +487,8 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch0_4_lat_dsl, ) # Case 3a - patch 1 - dreg_patch1_1_lon_vmask = np.where( - mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = np.where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) + dreg_patch1_1_lon_vmask = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) dreg_patch1_2_lon_vmask = np.where( mask_case3a, np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), @@ -543,12 +499,8 @@ def divide_flux_area_list_stencil_01_numpy( np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = np.where( - mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask - ) + dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) dreg_patch1_4_lon_vmask = np.where( mask_case3a, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), @@ -593,18 +545,10 @@ def divide_flux_area_list_stencil_01_numpy( tri_line2_p2_lat, ) # Case 3b - patch 0 - dreg_patch0_1_lon_dsl = np.where( - mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl - ) - dreg_patch0_1_lat_dsl = np.where( - mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl - ) - dreg_patch0_4_lon_dsl = np.where( - mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl - ) + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) dreg_patch0_2_lon_dsl = np.where( mask_case3b, np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), @@ -635,12 +579,8 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 - dreg_patch2_1_lon_vmask = np.where( - mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = np.where( - mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) + dreg_patch2_1_lon_vmask = np.where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) dreg_patch2_2_lon_vmask = np.where( mask_case3b, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), @@ -651,12 +591,8 @@ def divide_flux_area_list_stencil_01_numpy( np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = np.where( - mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask - ) + dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) dreg_patch2_4_lon_vmask = np.where( mask_case3b, np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), @@ -675,9 +611,7 @@ def divide_flux_area_list_stencil_01_numpy( [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] ) # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = np.logical_and.reduce( - [famask_bool, np.logical_not(indices_previously_matched)] - ) + mask_case4 = np.logical_and.reduce([famask_bool, np.logical_not(indices_previously_matched)]) # Case 4 - patch 0 - no change # Case 4 - patch 1 dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py index fceb4234db..7423360a29 100644 --- a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py @@ -19,11 +19,7 @@ divide_flux_area_list_stencil_02, ) from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - random_mask, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, random_mask from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py index 69412e0b7d..94ce83f87e 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py @@ -15,15 +15,9 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import ( - face_val_ppm_stencil_01, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - _shape, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -40,16 +34,10 @@ def face_val_ppm_stencil_01_numpy( zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] ) - zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / ( - p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1] - ) + zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) z_slope_a = ( p_cellhgt_mc_now[:, 1:-1] - / ( - p_cellhgt_mc_now[:, :-2] - + p_cellhgt_mc_now[:, 1:-1] - + p_cellhgt_mc_now[:, 2:] - ) + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) ) * ( (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 @@ -59,16 +47,10 @@ def face_val_ppm_stencil_01_numpy( zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] ) - zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1] - ) + zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) z_slope_b = ( p_cellhgt_mc_now[:, 1:-1] - / ( - p_cellhgt_mc_now[:, :-2] - + p_cellhgt_mc_now[:, 1:-1] - + p_cellhgt_mc_now[:, 1:-1] - ) + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) ) * ( (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py index 461c0bffd0..f00485f9b5 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py @@ -15,9 +15,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator import embedded as it_embedded -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import ( - face_val_ppm_stencil_02, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import _shape, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -36,14 +34,9 @@ def face_val_ppm_stencil_02_numpy( p_face_a = p_face_in - p_face_a[:, 1:] = p_cc[:, 1:] * ( - 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) - ) + ( + p_face_a[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) - ) * ( - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] - + p_cc[:, :-1] - ) + ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) p_face = np.where((vert_idx == slevp1) | (vert_idx == elev), p_face_a, p_face_in) p_face = np.where((vert_idx == slev), p_cc, p_face) diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py index a828e5897f..46364ca6c2 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import ( - face_val_ppm_stencil_02a, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -28,14 +26,9 @@ def face_val_ppm_stencil_02a_numpy( p_face = p_cc.copy() - p_face[:, 1:] = p_cc[:, 1:] * ( - 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) - ) + ( + p_face[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) - ) * ( - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] - + p_cc[:, :-1] - ) + ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) return p_face diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py index 6b1b867879..00a073f780 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import ( - face_val_ppm_stencil_02b, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py index a02c6d25d6..12191447d0 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import ( - face_val_ppm_stencil_02c, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py index 7dd8ca373f..cd5744078c 100644 --- a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import ( - face_val_ppm_stencil_05, -) +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py index 826c0fdfe1..df00eda2dd 100644 --- a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -18,11 +18,7 @@ hflux_ffsl_hybrid_stencil_01a, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py index 11417d34c4..7906fc3cff 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py @@ -18,10 +18,7 @@ hflx_limiter_mo_stencil_01b, ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py index 55a025afb5..a4a730509d 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -14,15 +14,9 @@ import numpy as np from numpy import int32 -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import ( - hflx_limiter_mo_stencil_02, -) +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import hflx_limiter_mo_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -45,12 +39,8 @@ def hflx_limiter_mo_stencil_02_numpy( np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), z_tracer_new_low, ) - z_tracer_max_out = np.where( - condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max - ) - z_tracer_min_out = np.where( - condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min - ) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) return z_tracer_new_out, z_tracer_max_out, z_tracer_min_out diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py index d2d29d7ca6..c2758faa62 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import ( - hflx_limiter_mo_stencil_04, -) +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py index 3533162d39..f837ac13b3 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py @@ -14,15 +14,9 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import ( - hflx_limiter_pd_stencil_01, -) +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import hflx_limiter_pd_stencil_01 from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py index b127080555..c935882673 100644 --- a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import ( - hflx_limiter_pd_stencil_02, -) +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import constant_field, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py index 458119d1b1..2cd9e9a2cd 100644 --- a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py @@ -14,14 +14,9 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import ( - hor_adv_stencil_01, -) +from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 4522ed5d9c..6f7cd10dfc 100644 --- a/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -71,12 +71,10 @@ def reference( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py index ed67d56456..1bdc3b7144 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -18,11 +18,7 @@ prep_gauss_quadrature_c_list_stencil, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -247,9 +243,7 @@ def prep_gauss_quadrature_c_list_stencil_numpy( + shape_func_4_4 * p_coords_dreg_v_4_y ) - p_quad_vector_sum_1 = ( - wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - ) + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 p_quad_vector_sum_2 = ( wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py index 315fbba080..a58cda650e 100644 --- a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py @@ -206,9 +206,7 @@ def prep_gauss_quadrature_c_stencil_numpy( + shape_func_4_4 * p_coords_dreg_v_4_y ) - p_quad_vector_sum_1 = ( - wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - ) + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 p_quad_vector_sum_2 = ( wgt_t_detjac_1 * z_gauss_pts_1_x + wgt_t_detjac_2 * z_gauss_pts_2_x diff --git a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py index 2bf78bc992..a66751bf99 100644 --- a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py +++ b/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import ( - rbf_intp_edge_stencil_01, -) +from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py index f7a50a179c..41ea35df79 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py @@ -16,15 +16,9 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.advection.recon_lsq_cell_c_stencil import ( - recon_lsq_cell_c_stencil, -) +from icon4py.model.atmosphere.advection.recon_lsq_cell_c_stencil import recon_lsq_cell_c_stencil from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -638,9 +632,7 @@ def test_recon_lsq_cell_c_stencil(): p_coeff_10, offset_provider={ "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), - "C2CECEC": StridedNeighborOffsetProvider( - CellDim, CECECDim, mesh.n_c2e2c2e2c - ), + "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), }, ) co1 = np.asarray(p_coeff_1) diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py index dc2d8f1203..74163d0ff2 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py @@ -18,11 +18,7 @@ recon_lsq_cell_c_svd_stencil, ) from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -317,9 +313,7 @@ def test_recon_lsq_cell_c_svd_stencil(): p_coeff_10_dsl, offset_provider={ "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), - "C2CECEC": StridedNeighborOffsetProvider( - CellDim, CECECDim, mesh.n_c2e2c2e2c - ), + "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), }, ) co1 = np.asarray(p_coeff_1_dsl) diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py index 3f7dca6ede..5d72af698e 100644 --- a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py @@ -18,11 +18,7 @@ recon_lsq_cell_l_svd_stencil, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - as_1D_sparse_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py index ce87c96606..97ab7f4889 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_01.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_01 import ( - step_advection_stencil_01, -) +from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -29,8 +27,7 @@ def step_advection_stencil_01_numpy( pd_time: float, ) -> np.ndarray: tmp = pd_time * ( - p_mflx_contra_v[:, 1:] * deepatmo_divzl - - p_mflx_contra_v[:, :-1] * deepatmo_divzu + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu ) return rhodz_ast + tmp diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py index fa15a23593..c6604dde9a 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_02.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_02 import ( - step_advection_stencil_02, -) +from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -29,10 +27,7 @@ def step_advection_stencil_02_numpy( pd_time: float, ) -> np.ndarray: - tmp = ( - p_mflx_contra_v[:, 1:] * deepatmo_divzl - - p_mflx_contra_v[:, :-1] * deepatmo_divzu - ) + tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py index 02a75ffef1..87082a32d0 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_03.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_03 import ( - step_advection_stencil_03, -) +from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py index 7a57b2adc4..dca7722d01 100644 --- a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py +++ b/model/atmosphere/advection/tests/test_step_advection_stencil_04.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.step_advection_stencil_04 import ( - step_advection_stencil_04, -) +from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py index 82b5b57d3e..79a0adc68e 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py @@ -18,11 +18,7 @@ upwind_hflux_miura3_stencil_01, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - random_field, - random_mask, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py index 2c574b4855..e7460b7e56 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -18,11 +18,7 @@ upwind_hflux_miura_cycl_stencil_01, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - random_field, - random_mask, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py index cae5ac6d30..823b0e0464 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -39,17 +39,13 @@ def upwind_hflux_miura_cycl_stencil_02_numpy( z_tracer_mflx_c2e = z_tracer_mflx[c2e] z_rhofluxdiv_c_out = ( - np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) - if nsub == int32(1) - else z_rhofluxdiv_c + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c ) z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out - z_tracer_new_dsl = ( - z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl - ) / z_rho_new_dsl + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py index db075b6b39..91470d8eb4 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -26,9 +26,7 @@ def upwind_hflux_miura_cycl_stencil_03b_numpy( z_tracer_mflx_2_dsl: np.array, z_tracer_mflx_3_dsl: np.array, ): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float( - 3 - ) + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) return p_out_e diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py index 8ff9f7f03e..78a0e595c5 100644 --- a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py @@ -18,11 +18,7 @@ upwind_hflux_miura_stencil_01, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - constant_field, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py index 86e49c9711..e66c64f369 100644 --- a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -41,9 +41,7 @@ def test_upwind_vflux_ppm_stencil_01(): np.asarray(z_face_up), np.asarray(z_face_down), np.asarray(p_cc) ) - upwind_vflux_ppm_stencil_01( - z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={} - ) + upwind_vflux_ppm_stencil_01(z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={}) assert np.allclose(ref_z_delta_q, z_delta_q) assert np.allclose(ref_z_a1, z_a1) diff --git a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py index 74da8a299f..15e3f247ed 100644 --- a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py @@ -13,9 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import ( - vert_adv_stencil_01, -) +from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -34,10 +32,7 @@ def vert_adv_stencil_01_numpy( tracer_new = ( tracer_now * rhodz_now + p_dtime - * ( - p_mflx_tracer_v[:, 1:] * deepatmo_divzl - - p_mflx_tracer_v[:, :-1] * deepatmo_divzu - ) + * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) ) / rhodz_new return tracer_new diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py index b9816e140e..d199edec97 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py @@ -14,9 +14,7 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import ( - v_limit_prbl_sm_stencil_01, -) +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import v_limit_prbl_sm_stencil_01 from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py index 6f77e35993..9c58230fd8 100644 --- a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py @@ -14,15 +14,9 @@ import numpy as np from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import ( - v_limit_prbl_sm_stencil_02, -) +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import v_limit_prbl_sm_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - random_field, - random_mask, - zero_field, -) +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py index 2d7ee1b965..333c21116d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py @@ -71,9 +71,7 @@ def _apply_diffusion_to_vn( vn, nudgezone_diff, ), - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v - ), + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) if limited_area else where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 63e140458c..294eb2ed19 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -50,9 +50,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdx, dwdy = where( int32(0) < vert_idx, - _calculate_horizontal_gradients_for_turbulence( - w_old, geofac_grg_x, geofac_grg_y - ), + _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) @@ -69,9 +67,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( & (vert_idx < nrdmax) & (interior_idx <= horz_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, area, z_nabla2_c - ), + _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), w, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py index aa76bed1b8..313e9ac500 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py @@ -27,9 +27,7 @@ def _apply_nabla2_and_nabla4_global_to_vn( diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py index 1710922bd0..c2f1074899 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py @@ -35,6 +35,4 @@ def apply_nabla2_to_vn_in_lateral_boundary( vn: Field[[EdgeDim, KDim], float], fac_bdydiff_v: float, ): - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn - ) + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py index 3483966bad..42128fd8ad 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py @@ -36,6 +36,4 @@ def apply_nabla2_to_w_in_upper_damping_layer( cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w - ) + _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py index 60c3126a24..1c686fe580 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py @@ -38,6 +38,4 @@ def calculate_diagnostics_for_turbulence( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _calculate_diagnostics_for_turbulence( - div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:]) - ) + _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:])) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 34c02a07d8..49eec5e7cf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -33,9 +33,7 @@ def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( theta_v, theta_ref_mc, thresh_tdiff ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d - ) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d) return kh_smag_e diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py index d0fff9756b..5a0a98d830 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py @@ -37,6 +37,4 @@ def calculate_horizontal_gradients_for_turbulence( dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], ): - _calculate_horizontal_gradients_for_turbulence( - w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy) - ) + _calculate_horizontal_gradients_for_turbulence(w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy)) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py index 65ac1c6238..cca6712de1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py @@ -39,6 +39,4 @@ def calculate_nabla2_for_theta( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _calculate_nabla2_for_theta( - kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp - ) + _calculate_nabla2_for_theta(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index ca00e06c58..37430b5448 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -32,6 +32,4 @@ def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: Field[[EdgeDim, KDim], float], enh_diffu_3d: Field[[CellDim, KDim], float], ): - _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d, out=kh_smag_e - ) + _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 06a4125d86..295be5cab3 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -33,6 +33,4 @@ def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( c_intp: Field[[VertexDim, V2CDim], float], p_vert_out: Field[[VertexDim, KDim], float], ): - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in, c_intp, out=p_vert_out - ) + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(p_cell_in, c_intp, out=p_vert_out) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py index 170414b2e6..0a9bf2d689 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -37,6 +37,4 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], ): - _mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out) - ) + _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py index 509a9e6150..ca293f5a91 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py @@ -24,9 +24,7 @@ def _mo_solve_nonhydro_stencil_02( exner_ref_mc: Field[[CellDim, KDim], float], exner_pr: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_exner_ex_pr = (1.0 + exner_exfac) * ( - exner - exner_ref_mc - ) - exner_exfac * exner_pr + z_exner_ex_pr = (1.0 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return z_exner_ex_pr, exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index bfb6b25f66..b709120573 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -43,9 +43,7 @@ def _mo_solve_nonhydro_stencil_10( ]: z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now(Koff[-1]) + wgt_nnew_rth * rho_var(Koff[-1]) - z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var( - Koff[-1] - ) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var(Koff[-1]) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 0818b0cc51..1247762663 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -33,13 +33,11 @@ def _compute_btraj( lvn_pos = where(p_vn > 0.0, True, False) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py index f7841c6d0b..4c628b66cf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py @@ -28,10 +28,7 @@ def _mo_solve_nonhydro_stencil_17( ) -> Field[[EdgeDim, KDim], float]: scalfac_dd3d = broadcast(scalfac_dd3d, (EdgeDim, KDim)) z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d - * scalfac_dd3d - * inv_dual_edge_length - * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) + hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) ) return z_graddiv_vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py index 9d99e03bab..5046f22bfa 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py @@ -23,9 +23,7 @@ def _mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = inv_dual_edge_length * ( - z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0]) - ) + z_gradh_exner = inv_dual_edge_length * (z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0])) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py index 737d0f57d2..8971bb8103 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py @@ -41,18 +41,12 @@ def _mo_solve_nonhydro_stencil_20( ( 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_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_dexner_dz_c1_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0])) ) ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py index f1919b1934..84fb050f6e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py @@ -25,9 +25,7 @@ def _mo_solve_nonhydro_stencil_22( z_hydro_corr: Field[[EdgeDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = where( - ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner - ) + z_gradh_exner = where(ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py index 66f3c22346..a23544d444 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py @@ -28,9 +28,7 @@ def _mo_solve_nonhydro_stencil_24( dtime: float, cpd: float, ) -> Field[[EdgeDim, KDim], float]: - vn_nnew = vn_nnow + dtime * ( - ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner - ) + vn_nnew = vn_nnow + dtime * (ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner) return vn_nnew diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index f50801788c..4b192e07b9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -25,9 +25,7 @@ def _mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: z_w_concorr_me_offset_1 = z_w_concorr_me(Koff[-1]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) z_w_concorr_mc_m0 = neighbor_sum(e_bln_c_s * z_w_concorr_me(C2E), axis=C2EDim) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 8c1f886763..5b57d8f19a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -28,15 +28,9 @@ def _mo_solve_nonhydro_stencil_40( z_w_concorr_me_offset_2 = z_w_concorr_me(Koff[-2]) z_w_concorr_me_offset_3 = z_w_concorr_me(Koff[-3]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) - z_w_concorr_mc_m2 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim - ) - z_w_concorr_mc_m3 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) + z_w_concorr_mc_m2 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim) + z_w_concorr_mc_m3 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim) return ( wgtfacq_c(Koff[-1]) * z_w_concorr_mc_m1 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py index 8606aa7198..5b6b0d51e2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py @@ -33,9 +33,7 @@ def _mo_solve_nonhydro_stencil_42( cpd: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c ) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) return z_w_expl, z_contr_w_fl_l diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py index d55620484a..a354730e7d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py @@ -42,17 +42,8 @@ def _mo_solve_nonhydro_stencil_55( rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( rho_ic * w - rho_ic(Koff[1]) * w(Koff[1]) ) - exner_new = ( - z_exner_expl - + exner_ref_mc - - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) - ) - theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new - ) + exner_new = z_exner_expl + exner_ref_mc - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) + theta_v_new = rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new return rho_new, exner_new, theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py index b1f718731a..e5e70e7ba2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py @@ -24,9 +24,7 @@ def _mo_solve_nonhydro_stencil_56_63( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - z_dwdz_dd = inv_ddqz_z_full * ( - (w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1])) - ) + z_dwdz_dd = inv_ddqz_z_full * ((w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1]))) return z_dwdz_dd diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py index 8b5347ada1..17879f2dbe 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py @@ -27,9 +27,7 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], float], r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py index d0435fdde2..995322d72d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py @@ -26,9 +26,7 @@ def _mo_solve_nonhydro_stencil_60( ndyn_substeps_var: float, dtime: float, ) -> Field[[CellDim, KDim], float]: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) return exner_dyn_incr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py index 3ff46a988b..eea826badc 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py @@ -30,9 +30,7 @@ def _mo_solve_nonhydro_stencil_65( r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py index 0c949bd699..6f46f314d2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py @@ -39,6 +39,4 @@ def mo_solve_nonhydro_stencil_67( rd_o_cvd: float, rd_o_p0ref: float, ): - _mo_solve_nonhydro_stencil_67( - rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner) - ) + _mo_solve_nonhydro_stencil_67(rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner)) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py index 8725c196c5..15230b8afb 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py @@ -31,10 +31,7 @@ def _mo_solve_nonhydro_stencil_68( ) -> Field[[CellDim, KDim], float]: theta_v_new = where( mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new, + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new, theta_v_new, ) return theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py index 72ad7e5471..c7c051da02 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py @@ -37,6 +37,4 @@ def mo_velocity_advection_stencil_04( vt: Field[[EdgeDim, KDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_04( - vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me - ) + _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py index 31e9161672..b4b4c697f8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py @@ -30,9 +30,7 @@ def _mo_velocity_advection_stencil_07( ) -> Field[[EdgeDim, KDim], float]: return vn_ie * inv_dual_edge_length * ( w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * ( - z_w_v(E2V[0]) - z_w_v(E2V[1]) - ) + ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * (z_w_v(E2V[0]) - z_w_v(E2V[1])) @program(grid_type=GridType.UNSTRUCTURED) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 4cdbc7d731..4dd324603b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -23,9 +23,7 @@ def _mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc( - Koff[-1] - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc(Koff[-1]) return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index d18b7afb77..b61a42a54a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py @@ -39,13 +39,9 @@ def _mo_velocity_advection_stencil_14( # should reduce the vertical cfl to a scalar and the levmask to a per level boolean field (see Fortran dycore). vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = where( - (cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c) - z_w_con_c = where( - (cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c) return cfl_clipping, vcfl, z_w_con_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py index 38b60dd96d..a1823ac6be 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py @@ -26,9 +26,7 @@ def _mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: ddt_w_adv = -z_w_con_c * ( - w(Koff[-1]) * coeff1_dwdz - - w(Koff[1]) * coeff2_dwdz - + w * (coeff2_dwdz - coeff1_dwdz) + w(Koff[-1]) * coeff1_dwdz - w(Koff[1]) * coeff2_dwdz + w * (coeff2_dwdz - coeff1_dwdz) ) return ddt_w_adv @@ -41,6 +39,4 @@ def mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ddt_w_adv: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_16( - z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:] - ) + _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py index 6476e2a57b..621e6befed 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py @@ -45,10 +45,7 @@ def _mo_velocity_advection_stencil_19( ): ddt_vn_adv = -( (coeff_gradekin(E2EC[0]) - coeff_gradekin(E2EC[1])) * z_kin_hor_e - + ( - -coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) - + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1]) - ) + + (-coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1])) + vt * (f_e + 0.5 * neighbor_sum(zeta(E2V), axis=E2VDim)) + neighbor_sum(z_w_con_c_full(E2C) * c_lin_e, axis=E2CDim) * (vn_ie - vn_ie(Koff[1])) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index c7d31cd5ef..fd189e6690 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -39,20 +39,11 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( 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]) - ) + * (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]) - ) + * (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]) - ) + * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) ) z_temp = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py index e495537311..bc53f68979 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py @@ -40,6 +40,4 @@ def update_theta_and_exner( exner: Field[[CellDim, KDim], float], rd_o_cvd: float, ): - _update_theta_and_exner( - z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner) - ) + _update_theta_and_exner(z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner)) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py index 18bc252014..d9d058adf3 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -44,12 +44,8 @@ def input_data(self, mesh): ) @staticmethod - def reference( - mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn - ): + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): area_edge = np.expand_dims(area_edge, axis=-1) diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py index 42af67e615..c3f4725214 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py @@ -34,9 +34,7 @@ def reference( ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum( - z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1 - ) + w = w - diff_multfac_w * area * area * np.sum(z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1) return dict(w=w) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py index fdb3fd9b81..c628918570 100644 --- a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py @@ -32,9 +32,7 @@ def reference( kc_offset_1 = np.roll(kh_c, shift=1, axis=1) div_offset_1 = np.roll(div, shift=1, axis=1) div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] - hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[ - :, 1: - ] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] return dict(div_ic=div_ic, hdef_ic=hdef_ic) @pytest.fixture @@ -44,6 +42,4 @@ def input_data(self, mesh): kh_c = random_field(mesh, CellDim, KDim) div_ic = zero_field(mesh, CellDim, KDim) hdef_ic = zero_field(mesh, CellDim, KDim) - return dict( - wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic - ) + return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index e39aa7c1b4..184dba0dca 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -111,9 +111,7 @@ def reference( + v_vert_e2c2v[:, 3] * primal_normal_vert_y[:, 3] ) - kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - ( - dvt_tang * inv_primal_edge_length - ) + kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - (dvt_tang * inv_primal_edge_length) kh_smag_2 = kh_smag_2 * kh_smag_2 diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py index 3a65230bd7..5b1f998750 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py @@ -29,14 +29,10 @@ class TestCalculateNabla2OfTheta(StencilTest): OUTPUTS = ("z_temp",) @staticmethod - def reference( - mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs) -> np.array: geofac_div = geofac_div.reshape(mesh.c2e.shape) geofac_div = np.expand_dims(geofac_div, axis=-1) - z_temp = np.sum( - z_nabla2_e[mesh.c2e] * geofac_div, axis=1 - ) # sum along edge dimension + z_temp = np.sum(z_nabla2_e[mesh.c2e] * geofac_div, axis=1) # sum along edge dimension return dict(z_temp=z_temp) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py index 7501aa0738..dab2dcb907 100644 --- a/model/atmosphere/dycore/tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/test_compute_airmass.py @@ -36,7 +36,5 @@ def test_compute_airmass(): ref = compute_airmass_numpy( np.asarray(rho_in), np.asarray(ddqz_z_full_in), np.asarray(deepatmo_t1mc_in) ) - compute_airmass( - rho_in, ddqz_z_full_in, deepatmo_t1mc_in, airmass_out, offset_provider={} - ) + compute_airmass(rho_in, ddqz_z_full_in, deepatmo_t1mc_in, airmass_out, offset_provider={}) assert np.allclose(airmass_out, ref) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py index c04d712ff5..05975778c4 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py @@ -34,9 +34,7 @@ def reference( exner_exfac: np.array, **kwargs, ) -> dict: - z_exner_ex_pr = (1 + exner_exfac) * ( - exner - exner_ref_mc - ) - exner_exfac * exner_pr + z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py index 5b1cf61da9..e247e9f7a9 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py @@ -33,12 +33,9 @@ def reference( z_exner_ic: np.array, ) -> np.array: z_exner_ic[:, 3:] = ( - np.roll(wgtfacq_c, shift=1, axis=1) - * np.roll(z_exner_ex_pr, shift=1, axis=1) - + np.roll(wgtfacq_c, shift=2, axis=1) - * np.roll(z_exner_ex_pr, shift=2, axis=1) - + np.roll(wgtfacq_c, shift=3, axis=1) - * np.roll(z_exner_ex_pr, shift=3, axis=1) + np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) + + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) + + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) )[:, 3:] return {"z_exner_ic": z_exner_ic} diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py index 7a5289af18..6d54d033aa 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py @@ -26,13 +26,9 @@ class TestMoSolveNonhydroStencil05(StencilTest): OUTPUTS = ("z_exner_ic",) @staticmethod - def reference( - mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs - ) -> np.array: + def reference(mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs) -> np.array: z_exner_ex_pr_offset_1 = np.roll(z_exner_ex_pr, shift=1, axis=1) - z_exner_ic = ( - wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 - ) + z_exner_ic = wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 return dict(z_exner_ic=z_exner_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py index 12ef9a3cb4..e6086dc445 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil06(StencilTest): OUTPUTS = ("z_dexner_dz_c_1",) @staticmethod - def reference( - mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs) -> np.array: z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full return dict(z_dexner_dz_c_1=z_dexner_dz_c_1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py index c2d1d80e2b..d66d399fa2 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py @@ -55,9 +55,7 @@ def reference( z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now_offset + wgt_nnew_rth * rho_var_offset - z_theta_tavg_m1 = ( - wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset - ) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( @@ -67,9 +65,7 @@ def reference( ) z_theta_v_pr_mc_m1 = z_theta_tavg_m1 - theta_ref_mc_offset z_theta_v_pr_mc = z_theta_tavg - theta_ref_mc - z_theta_v_pr_ic = ( - wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 - ) + z_theta_v_pr_ic = wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 theta_v_ic = ( wgtfac_c * z_theta_tavg + (1 - wgtfac_c) * z_theta_tavg_m1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py index 40f10d2032..bb92f1fbef 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil15(StencilTest): OUTPUTS = ("z_rho_e", "z_theta_v_e") @staticmethod - def reference( - mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs - ) -> tuple[np.array]: + def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> tuple[np.array]: z_rho_e = np.zeros_like(z_rho_e) z_theta_v_e = np.zeros_like(z_theta_v_e) return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 8dcbd9b884..461c0a654c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -48,12 +48,10 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py index 56d05c7b98..5a782f7c5f 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py @@ -41,9 +41,8 @@ def reference( z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - z_gradh_exner = ( - inv_dual_edge_length * z_exner_ex_weighted - - ddxn_z_full * np.sum(c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1) + z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( + c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1 ) return dict(z_gradh_exner=z_gradh_exner) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py index 6fa501ec87..320fb7dc4a 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py @@ -47,12 +47,8 @@ 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, :] * ( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py index c5d3166620..a601630e64 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py @@ -57,9 +57,7 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): full_shape, theta_v_ic, e2c, ikoffset ) - inv_ddqz_z_full_at_kidx, _ = _apply_index_field( - full_shape, inv_ddqz_z_full, e2c, ikoffset - ) + inv_ddqz_z_full_at_kidx, _ = _apply_index_field(full_shape, inv_ddqz_z_full, e2c, ikoffset) z_theta1 = ( theta_v_at_kidx[:, 0, :] diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py index fd5765f805..85fdb0788e 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil25(StencilTest): OUTPUTS = ("z_graddiv2_vn",) @staticmethod - def reference( - mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs - ) -> np.array: + def reference(mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs) -> np.array: geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) z_graddiv2_vn = np.sum(z_graddiv_vn[mesh.e2c2eO] * geofac_grdiv, axis=1) return dict(z_graddiv2_vn=z_graddiv2_vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py index dded703d10..ec2389a8de 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil26(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference( - mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs - ) -> dict: + def reference(mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs) -> dict: vn = vn + (scal_divdamp_o2 * z_graddiv_vn) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py index 2eeba98f41..17234d448c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil28(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference( - mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs - ) -> np.array: + def reference(mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs) -> np.array: vn = vn + (iau_wgt_dyn * vn_incr) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py index 1520b78be4..cc48d754d4 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil29(StencilTest): OUTPUTS = ("vn_new",) @staticmethod - def reference( - mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs - ) -> dict: + def reference(mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs) -> dict: vn_new = vn_now + dtime * grf_tend_vn return dict(vn_new=vn_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index 46062545d9..0530136399 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -36,12 +36,8 @@ def reference( e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_mc_m0 = np.sum(e_bln_c_s * z_w_concorr_me[mesh.c2e], axis=1) - z_w_concorr_mc_m1 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 - ) - w_concorr_c = ( - wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 - ) + z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) + w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index 896b38a04c..6774dc9461 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -37,15 +37,9 @@ def reference( z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_me_offset_2 = np.roll(z_w_concorr_me, shift=2, axis=1) z_w_concorr_me_offset_3 = np.roll(z_w_concorr_me, shift=3, axis=1) - z_w_concorr_mc_m1 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 - ) - z_w_concorr_mc_m2 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1 - ) - z_w_concorr_mc_m3 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1 - ) + z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) + z_w_concorr_mc_m2 = np.sum(e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1) + z_w_concorr_mc_m3 = np.sum(e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1) w_concorr_c = ( np.roll(wgtfacq_c, shift=1, axis=1) * z_w_concorr_mc_m1 + np.roll(wgtfacq_c, shift=2, axis=1) * z_w_concorr_mc_m2 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py index 6835c5fbde..2e0882c12f 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py @@ -42,9 +42,7 @@ def reference( **kwargs, ) -> tuple[np.array]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c ) vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py index ea0e2aed07..6008665b5e 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py @@ -40,9 +40,7 @@ def reference( cvd, **kwargs, ) -> dict: - z_beta = ( - dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full - ) + z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py index 5892ce5b6d..66aa9b6b09 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil47(StencilTest): OUTPUTS = ("w_nnew", "z_contr_w_fl_l") @staticmethod - def reference( - mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs - ) -> dict: + def reference(mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: w_nnew = w_concorr_c z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py index d1cb07a661..f7e12c96f2 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py @@ -57,9 +57,7 @@ def mo_solve_nonhydro_stencil_51_numpy( z_c = -z_gamma * z_beta * z_alpha_k_plus_1 z_b = 1.0 + z_gamma * z_alpha[:, :-1] * (z_beta_k_minus_1 + z_beta) z_q = mo_solve_nonhydro_stencil_51_z_q_numpy(z_c, z_b) - w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy( - z_gamma, z_b, z_w_expl, z_exner_expl - ) + w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy(z_gamma, z_b, z_w_expl, z_exner_expl) return z_q, w_nnew diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py index 85b1531896..6c0cc28e0d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py @@ -49,15 +49,11 @@ def mo_solve_nonhydro_stencil_52_numpy( for k in range(1, k_size): z_a[:, k] = -z_gamma[:, k] * z_beta[:, k - 1] * z_alpha[:, k - 1] z_c[:, k] = -z_gamma[:, k] * z_beta[:, k] * z_alpha[:, k + 1] - z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * ( - z_beta[:, k - 1] + z_beta[:, k] - ) + z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * (z_beta[:, k - 1] + z_beta[:, k]) z_g[:, k] = 1.0 / (z_b[:, k] + z_a[:, k] * z_q[:, k - 1]) z_q[:, k] = -z_c[:, k] * z_g[:, k] - w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * ( - z_exner_expl[:, k - 1] - z_exner_expl[:, k] - ) + w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * (z_exner_expl[:, k - 1] - z_exner_expl[:, k]) w[:, k] = (w[:, k] - z_a[:, k] * w[:, k - 1]) * z_g[:, k] return z_q, w diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py index bc2e501a21..eeb17c2cce 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil54(StencilTest): OUTPUTS = ("w",) @staticmethod - def reference( - mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -> np.array: z_raylfac = np.expand_dims(z_raylfac, axis=0) w_1 = np.expand_dims(w_1, axis=-1) w = z_raylfac * w + (1.0 - z_raylfac) * w_1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py index 1ba4083069..b8e2166397 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py @@ -58,10 +58,7 @@ def reference( - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) ) theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new ) return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py index 93e8c6a95d..dff001052b 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py @@ -37,9 +37,7 @@ def reference( **kwargs, ) -> dict: vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return dict(mass_flx_ic=mass_flx_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py index 407dc06e90..634446cb33 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py @@ -35,9 +35,7 @@ def reference( dtime: float, **kwargs, ) -> np.array: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) return dict(exner_dyn_incr=exner_dyn_incr) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py index d31fffed0a..9c3d3442cf 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil62(StencilTest): OUTPUTS = ("w_new",) @staticmethod - def reference( - mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs - ) -> np.array: + def reference(mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs) -> np.array: w_new = w_now + dtime * grf_tend_w return dict(w_new=w_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py index 77e46f5927..b5b5b9c0c0 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py @@ -41,9 +41,7 @@ def reference( vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return dict(mass_flx_ic=mass_flx_ic) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py index 081dbb08d9..3a53d52702 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py @@ -42,10 +42,7 @@ def reference( theta_v_new = np.where( mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) - / rho_new, + rho_now * theta_v_now * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) / rho_new, theta_v_new, ) return dict(theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index 60fabe0600..142b7e56cb 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -26,24 +26,18 @@ class TestMoVelocityAdvectionStencil02VnIe(StencilTest): OUTPUTS = ("vn_ie", "z_kin_hor_e") @staticmethod - def mo_velocity_advection_stencil_02_vn_ie_numpy( - wgtfac_e: np.array, vn: np.array - ) -> np.array: + def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie @staticmethod - def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( - vn: np.array, vt: np.array - ) -> np.array: + def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e @classmethod - def reference( - cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs - ) -> dict: + def reference(cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs) -> dict: vn_ie = cls.mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e, vn) z_kin_hor_e = cls.mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn, vt) return dict(vn_ie=vn_ie, z_kin_hor_e=z_kin_hor_e) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py index 0a8bf8d0f1..562129855b 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil08(StencilTest): OUTPUTS = ("z_ekinh",) @staticmethod - def reference( - mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_ekinh = np.sum(z_kin_hor_e[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_ekinh=z_ekinh) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py index 1b0d8d4f1a..e120054230 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil09(StencilTest): OUTPUTS = ("z_w_concorr_mc",) @staticmethod - def reference( - mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_mc = np.sum(z_w_concorr_me[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_w_concorr_mc=z_w_concorr_mc) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index c561da5a82..d7d7e3d2c7 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -26,13 +26,9 @@ class TestMoVelocityAdvectionStencil10(StencilTest): OUTPUTS = ("w_concorr_c",) @staticmethod - def reference( - mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs - ) -> np.array: + def reference(mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = ( - wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py index 336ce24d35..eb7fb3fc44 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil13(StencilTest): OUTPUTS = ("z_w_con_c",) @staticmethod - def reference( - mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs - ) -> np.array: + def reference(mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> np.array: z_w_con_c = z_w_con_c - w_concorr_c return dict(z_w_con_c=z_w_con_c) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py index 330b4c9b0a..88063c317b 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py @@ -55,10 +55,7 @@ def mo_velocity_advection_stencil_19_numpy( ddt_vn_adv = -( (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + ( - -coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] - + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1] - ) + + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + vt * (f_e + 0.5 * np.sum(zeta[e2v], axis=1)) + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) * (vn_ie[:, :-1] - vn_ie[:, 1:]) diff --git a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 00b8ffce42..b0114cc26a 100644 --- a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -56,15 +56,12 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( ] sum_over = np.sum( - geofac_n2s_nbh - * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), axis=1, ) geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - return np.where( - mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp - ) + return np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) def test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points(): From 385168754bedd6648f0becc01c67087ed4f131a0 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 8 Sep 2023 17:43:54 +0200 Subject: [PATCH 094/105] Change names for test folders to '_tests' and add __init__.py --- .../atmosphere/advection/advection_tests/__init__.py | 12 ++++++++++++ .../advection/{tests => advection_tests}/conftest.py | 0 .../test_btraj_dreg_stencil_01.py | 0 .../test_btraj_dreg_stencil_02.py | 0 .../test_btraj_dreg_stencil_03.py | 0 .../test_divide_flux_area_list_stencil_01.py | 0 .../test_divide_flux_area_list_stencil_02.py | 0 .../test_face_val_ppm_stencil_01.py | 0 .../test_face_val_ppm_stencil_02.py | 0 .../test_face_val_ppm_stencil_02a.py | 0 .../test_face_val_ppm_stencil_02b.py | 0 .../test_face_val_ppm_stencil_02c.py | 0 .../test_face_val_ppm_stencil_05.py | 0 .../test_hflux_ffsl_hybrid_stencil_01a.py | 0 .../test_hflux_ffsl_hybrid_stencil_02.py | 0 .../test_hflx_limiter_mo_stencil_01a.py | 0 .../test_hflx_limiter_mo_stencil_01b.py | 0 .../test_hflx_limiter_mo_stencil_02.py | 0 .../test_hflx_limiter_mo_stencil_03.py | 0 .../test_hflx_limiter_mo_stencil_04.py | 0 .../test_hflx_limiter_pd_stencil_01.py | 0 .../test_hflx_limiter_pd_stencil_02.py | 0 .../test_hor_adv_stencil_01.py | 0 .../test_mo_advection_traj_btraj_compute_o1_dsl.py | 0 .../test_prep_gauss_quadrature_c_list_stencil.py | 0 .../test_prep_gauss_quadrature_c_stencil.py | 0 .../test_rbf_intp_edge_stencil_01.py | 0 .../test_recon_lsq_cell_c_stencil.py | 0 .../test_recon_lsq_cell_c_svd_stencil.py | 0 .../test_recon_lsq_cell_l_svd_stencil.py | 0 .../{tests => advection_tests}/test_set_zero_c.py | 0 .../{tests => advection_tests}/test_set_zero_c_k.py | 0 .../test_step_advection_stencil_01.py | 0 .../test_step_advection_stencil_02.py | 0 .../test_step_advection_stencil_03.py | 0 .../test_step_advection_stencil_04.py | 0 .../test_upwind_hflux_miura3_stencil_01.py | 0 .../test_upwind_hflux_miura_cycl_stencil_01.py | 0 .../test_upwind_hflux_miura_cycl_stencil_02.py | 0 .../test_upwind_hflux_miura_cycl_stencil_03a.py | 0 .../test_upwind_hflux_miura_cycl_stencil_03b.py | 0 .../test_upwind_hflux_miura_stencil_01.py | 0 .../test_upwind_vflux_ppm_stencil_01.py | 0 .../test_vert_adv_stencil_01.py | 0 .../test_vlimit_prbl_sm_stencil_01.py | 0 .../test_vlimit_prbl_sm_stencil_02.py | 0 model/atmosphere/dycore/tests/__init__.py | 12 ++++++++++++ 47 files changed, 24 insertions(+) create mode 100644 model/atmosphere/advection/advection_tests/__init__.py rename model/atmosphere/advection/{tests => advection_tests}/conftest.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_btraj_dreg_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_btraj_dreg_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_btraj_dreg_stencil_03.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_divide_flux_area_list_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_divide_flux_area_list_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_02a.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_02b.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_02c.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_face_val_ppm_stencil_05.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflux_ffsl_hybrid_stencil_01a.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflux_ffsl_hybrid_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_mo_stencil_01a.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_mo_stencil_01b.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_mo_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_mo_stencil_03.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_mo_stencil_04.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_pd_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hflx_limiter_pd_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_hor_adv_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_mo_advection_traj_btraj_compute_o1_dsl.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_prep_gauss_quadrature_c_list_stencil.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_prep_gauss_quadrature_c_stencil.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_rbf_intp_edge_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_recon_lsq_cell_c_stencil.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_recon_lsq_cell_c_svd_stencil.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_recon_lsq_cell_l_svd_stencil.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_set_zero_c.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_set_zero_c_k.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_step_advection_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_step_advection_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_step_advection_stencil_03.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_step_advection_stencil_04.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura3_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura_cycl_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura_cycl_stencil_02.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura_cycl_stencil_03a.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura_cycl_stencil_03b.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_hflux_miura_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_upwind_vflux_ppm_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_vert_adv_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_vlimit_prbl_sm_stencil_01.py (100%) rename model/atmosphere/advection/{tests => advection_tests}/test_vlimit_prbl_sm_stencil_02.py (100%) create mode 100644 model/atmosphere/dycore/tests/__init__.py diff --git a/model/atmosphere/advection/advection_tests/__init__.py b/model/atmosphere/advection/advection_tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/atmosphere/advection/tests/conftest.py b/model/atmosphere/advection/advection_tests/conftest.py similarity index 100% rename from model/atmosphere/advection/tests/conftest.py rename to model/atmosphere/advection/advection_tests/conftest.py diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_btraj_dreg_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_btraj_dreg_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/test_btraj_dreg_stencil_03.py rename to model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_divide_flux_area_list_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_02a.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_02b.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_02c.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py diff --git a/model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py similarity index 100% rename from model/atmosphere/advection/tests/test_face_val_ppm_stencil_05.py rename to model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_01a.py rename to model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py diff --git a/model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflux_ffsl_hybrid_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01a.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_01b.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_03.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_mo_stencil_04.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_hflx_limiter_pd_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_hor_adv_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/advection_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 100% rename from model/atmosphere/advection/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/advection_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_list_stencil.py rename to model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py diff --git a/model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/test_prep_gauss_quadrature_c_stencil.py rename to model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py diff --git a/model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_rbf_intp_edge_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/test_recon_lsq_cell_c_stencil.py rename to model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/test_recon_lsq_cell_c_svd_stencil.py rename to model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py diff --git a/model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/test_recon_lsq_cell_l_svd_stencil.py rename to model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.py diff --git a/model/atmosphere/advection/tests/test_set_zero_c.py b/model/atmosphere/advection/advection_tests/test_set_zero_c.py similarity index 100% rename from model/atmosphere/advection/tests/test_set_zero_c.py rename to model/atmosphere/advection/advection_tests/test_set_zero_c.py diff --git a/model/atmosphere/advection/tests/test_set_zero_c_k.py b/model/atmosphere/advection/advection_tests/test_set_zero_c_k.py similarity index 100% rename from model/atmosphere/advection/tests/test_set_zero_c_k.py rename to model/atmosphere/advection/advection_tests/test_set_zero_c_k.py diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_step_advection_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_step_advection_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/test_step_advection_stencil_03.py rename to model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py diff --git a/model/atmosphere/advection/tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py similarity index 100% rename from model/atmosphere/advection/tests/test_step_advection_stencil_04.py rename to model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura3_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03a.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura_cycl_stencil_03b.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py diff --git a/model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_hflux_miura_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_upwind_vflux_ppm_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_vert_adv_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_01.py rename to model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py diff --git a/model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/test_vlimit_prbl_sm_stencil_02.py rename to model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py diff --git a/model/atmosphere/dycore/tests/__init__.py b/model/atmosphere/dycore/tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/dycore/tests/__init__.py @@ -0,0 +1,12 @@ +# 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 a95c8206703fdc869cb706bcdee1744362d6251a Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 12 Sep 2023 16:04:23 +0200 Subject: [PATCH 095/105] Merge greenline to main (#260) merging of greenline with main --- .github/workflows/icon4py-qa.yml | 55 +- .github/workflows/icon4pytools-qa.yml | 4 + .gitignore | 2 + base-requirements-dev.txt | 5 + base-requirements.txt | 1 + model/README.md | 46 + model/atmosphere/diffusion/.bumpversion.cfg | 10 + model/atmosphere/diffusion/.flake8 | 42 + .../diffusion/.pre-commit-config.yaml | 114 +++ model/atmosphere/diffusion/README.md | 9 + .../diffusion/diffusion_tests/__init__.py | 12 + .../diffusion/diffusion_tests/conftest.py | 94 ++ .../diffusion_tests/mpi_tests/__init__.py | 12 + .../mpi_tests/test_parallel_diffusion.py | 110 +++ ...st_apply_nabla2_and_nabla4_global_to_vn.py | 2 +- .../test_apply_nabla2_and_nabla4_to_vn.py | 36 +- ..._apply_nabla2_to_vn_in_lateral_boundary.py | 2 +- .../test_apply_nabla2_to_w.py | 14 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 8 +- ...st_calculate_diagnostics_for_turbulence.py | 2 +- ...ate_horizontal_gradients_for_turbulence.py | 7 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../test_calculate_nabla2_for_w.py | 9 +- .../test_calculate_nabla2_for_z.py | 4 +- .../test_calculate_nabla2_of_theta.py | 4 +- .../diffusion_tests}/test_calculate_nabla4.py | 2 +- .../diffusion_tests/test_diffusion.py | 348 +++++++ .../diffusion_tests/test_diffusion_states.py | 28 + .../diffusion_tests/test_diffusion_utils.py | 147 +++ ...n_coefficient_for_grid_point_cold_pools.py | 2 +- ...d_for_grid_point_cold_pools_enhancement.py | 2 +- ...orary_fields_for_turbulence_diagnostics.py | 20 +- ...fusion_nabla_of_theta_over_steep_points.py | 6 +- .../test_update_theta_and_exner.py | 9 +- .../diffusion/diffusion_tests/utils.py | 77 ++ model/atmosphere/diffusion/pyproject.toml | 120 +++ .../atmosphere/diffusion/requirements-dev.txt | 3 + model/atmosphere/diffusion/requirements.txt | 3 + .../model/atmosphere/diffusion/__init__.py | 34 + .../model/atmosphere/diffusion/diffusion.py | 850 ++++++++++++++++++ .../atmosphere/diffusion/diffusion_states.py | 115 +++ .../atmosphere/diffusion/diffusion_utils.py | 230 +++++ .../atmosphere/diffusion/stencils/__init__.py | 12 + .../apply_diffusion_to_theta_and_exner.py | 14 +- .../stencils}/apply_diffusion_to_vn.py | 20 +- ...ute_horizontal_gradients_for_turbulance.py | 18 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 0 .../apply_nabla2_and_nabla4_to_vn.py | 0 .../apply_nabla2_to_vn_in_lateral_boundary.py | 0 .../diffusion/stencils}/apply_nabla2_to_w.py | 19 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 18 +- ...te_diagnostic_quantities_for_turbulence.py | 24 +- .../calculate_diagnostics_for_turbulence.py | 0 ..._coefficients_for_grid_point_cold_pools.py | 20 +- ...ate_horizontal_gradients_for_turbulence.py | 17 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 0 .../stencils}/calculate_nabla2_for_theta.py | 26 +- .../stencils}/calculate_nabla2_for_w.py | 16 +- .../stencils}/calculate_nabla2_for_z.py | 0 .../stencils}/calculate_nabla2_of_theta.py | 0 .../diffusion/stencils}/calculate_nabla4.py | 0 ...n_coefficient_for_grid_point_cold_pools.py | 0 ...d_for_grid_point_cold_pools_enhancement.py | 0 ...orary_fields_for_turbulence_diagnostics.py | 14 +- ...fusion_nabla_of_theta_over_steep_points.py | 8 + .../stencils}/update_theta_and_exner.py | 19 +- model/atmosphere/diffusion/tests/conftest.py | 24 + model/atmosphere/dycore/pyproject.toml | 2 +- model/atmosphere/dycore/tests/conftest.py | 35 +- model/common/.pre-commit-config.yaml | 10 + model/common/pyproject.toml | 5 +- .../src/icon4py/model/common/constants.py | 31 + .../model/common/decomposition/__init__.py | 12 + .../model/common/decomposition/decomposed.py | 251 ++++++ .../common/decomposition/parallel_setup.py | 98 ++ .../src/icon4py/model/common/dimension.py | 3 + .../src/icon4py/model/common/grid/__init__.py | 12 + .../icon4py/model/common/grid/grid_manager.py | 405 +++++++++ .../icon4py/model/common/grid/horizontal.py | 269 ++++++ .../icon4py/model/common/grid/icon_grid.py | 197 ++++ .../src/icon4py/model/common/grid/vertical.py | 56 ++ .../model/common/interpolation/__init__.py | 12 + .../interpolation/interpolation_fields.py | 51 ++ .../common/interpolation/stencils/__init__.py | 12 + .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 17 +- .../model/common/test_utils/data_handling.py | 30 + .../model/common/test_utils/fixtures.py | 183 ++++ .../common/test_utils/parallel_helpers.py | 22 + .../model/common/test_utils/pytest_config.py | 40 + .../common/test_utils/serialbox_utils.py | 577 ++++++++++++ .../model/common/test_utils/simple_mesh.py | 39 +- .../common/src/icon4py/model/common/utils.py | 22 + model/common/tests/conftest.py | 64 ++ .../common/tests/mpi_tests/test_decomposed.py | 196 ++++ .../tests/mpi_tests/test_parallel_setup.py | 47 + model/common/tests/test_grid_manager.py | 521 +++++++++++ model/common/tests/test_icon_grid.py | 400 +++++++++ .../common/tests/test_interpolation_fields.py | 51 ++ ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 9 +- model/common/tests/test_vertical.py | 47 + model/driver/.bumpversion.cfg | 10 + model/driver/.flake8 | 42 + model/driver/.pre-commit-config.yaml | 114 +++ model/driver/README.md | 36 + model/driver/pyproject.toml | 126 +++ model/driver/requirements-dev.txt | 5 + model/driver/requirements.txt | 5 + .../src/icon4py/model/driver/__init__.py | 33 + .../src/icon4py/model/driver/dycore_driver.py | 295 ++++++ .../model/driver/icon_configuration.py | 85 ++ .../src/icon4py/model/driver/io_utils.py | 187 ++++ model/driver/tests/conftest.py | 24 + model/driver/tests/test_io_utils.py | 108 +++ model/requirements-dev.txt | 2 + model/requirements.txt | 2 + model/tox.ini | 3 +- requirements-dev.txt | 2 + requirements.txt | 2 + tools/README.md | 36 + tools/pyproject.toml | 14 +- tools/requirements-dev.txt | 1 + .../src/icon4pytools/liskov/external/gt4py.py | 6 +- tools/src/icon4pytools/py2f/__init__.py | 12 + tools/src/icon4pytools/py2f/cffi_utils.py | 177 ++++ tools/src/icon4pytools/py2f/codegen.py | 145 +++ tools/src/icon4pytools/py2f/parsing.py | 43 + tools/src/icon4pytools/py2f/py2fgen.py | 51 ++ tools/src/icon4pytools/py2f/typing_utils.py | 27 + .../icon4pytools/py2f/wrappers/__init__.py | 12 + .../py2f/wrappers/diffusion_wrapper.py | 265 ++++++ tools/tests/icon4pygen/test_codegen.py | 26 +- tools/tests/liskov/test_external.py | 2 +- tools/tests/py2f/test_cffi_utils.py | 80 ++ tools/tests/py2f/test_code_generation.py | 140 +++ tools/tests/py2f/test_parsing_wrapper.py | 39 + tools/tests/py2f/test_py2fgen.py | 25 + tox.ini | 2 +- 137 files changed, 8628 insertions(+), 153 deletions(-) create mode 100644 model/atmosphere/diffusion/.bumpversion.cfg create mode 100644 model/atmosphere/diffusion/.flake8 create mode 100644 model/atmosphere/diffusion/.pre-commit-config.yaml create mode 100644 model/atmosphere/diffusion/README.md create mode 100644 model/atmosphere/diffusion/diffusion_tests/__init__.py create mode 100644 model/atmosphere/diffusion/diffusion_tests/conftest.py create mode 100644 model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py create mode 100644 model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_apply_nabla2_and_nabla4_global_to_vn.py (95%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_apply_nabla2_and_nabla4_to_vn.py (62%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_apply_nabla2_to_vn_in_lateral_boundary.py (94%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_apply_nabla2_to_w.py (80%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_apply_nabla2_to_w_in_upper_damping_layer.py (83%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_diagnostics_for_turbulence.py (94%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_horizontal_gradients_for_turbulence.py (85%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_nabla2_and_smag_coefficients_for_vn.py (98%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_nabla2_for_w.py (81%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_nabla2_for_z.py (94%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_nabla2_of_theta.py (93%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_calculate_nabla4.py (97%) create mode 100644 model/atmosphere/diffusion/diffusion_tests/test_diffusion.py create mode 100644 model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py create mode 100644 model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py (92%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_temporary_field_for_grid_point_cold_pools_enhancement.py (94%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_temporary_fields_for_turbulence_diagnostics.py (74%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (93%) rename model/atmosphere/{dycore/tests => diffusion/diffusion_tests}/test_update_theta_and_exner.py (84%) create mode 100644 model/atmosphere/diffusion/diffusion_tests/utils.py create mode 100644 model/atmosphere/diffusion/pyproject.toml create mode 100644 model/atmosphere/diffusion/requirements-dev.txt create mode 100644 model/atmosphere/diffusion/requirements.txt create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_theta_and_exner.py (86%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_vn.py (85%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py (83%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_and_nabla4_global_to_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_and_nabla4_to_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_vn_in_lateral_boundary.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_w.py (75%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_w_in_upper_damping_layer.py (74%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_diagnostic_quantities_for_turbulence.py (71%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_diagnostics_for_turbulence.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py (72%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_horizontal_gradients_for_turbulence.py (76%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_and_smag_coefficients_for_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_theta.py (67%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_w.py (74%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_z.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_of_theta.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla4.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/enhance_diffusion_coefficient_for_grid_point_cold_pools.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/temporary_field_for_grid_point_cold_pools_enhancement.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/temporary_fields_for_turbulence_diagnostics.py (77%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (92%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/update_theta_and_exner.py (76%) create mode 100644 model/atmosphere/diffusion/tests/conftest.py create mode 100644 model/common/src/icon4py/model/common/constants.py create mode 100644 model/common/src/icon4py/model/common/decomposition/__init__.py create mode 100644 model/common/src/icon4py/model/common/decomposition/decomposed.py create mode 100644 model/common/src/icon4py/model/common/decomposition/parallel_setup.py create mode 100644 model/common/src/icon4py/model/common/grid/__init__.py create mode 100644 model/common/src/icon4py/model/common/grid/grid_manager.py create mode 100644 model/common/src/icon4py/model/common/grid/horizontal.py create mode 100644 model/common/src/icon4py/model/common/grid/icon_grid.py create mode 100644 model/common/src/icon4py/model/common/grid/vertical.py create mode 100644 model/common/src/icon4py/model/common/interpolation/__init__.py create mode 100644 model/common/src/icon4py/model/common/interpolation/interpolation_fields.py create mode 100644 model/common/src/icon4py/model/common/interpolation/stencils/__init__.py rename model/{atmosphere/dycore/src/icon4py/model/atmosphere/dycore => common/src/icon4py/model/common/interpolation/stencils}/mo_intp_rbf_rbf_vec_interpol_vertex.py (76%) create mode 100644 model/common/src/icon4py/model/common/test_utils/data_handling.py create mode 100644 model/common/src/icon4py/model/common/test_utils/fixtures.py create mode 100644 model/common/src/icon4py/model/common/test_utils/parallel_helpers.py create mode 100644 model/common/src/icon4py/model/common/test_utils/pytest_config.py create mode 100644 model/common/src/icon4py/model/common/test_utils/serialbox_utils.py create mode 100644 model/common/src/icon4py/model/common/utils.py create mode 100644 model/common/tests/conftest.py create mode 100644 model/common/tests/mpi_tests/test_decomposed.py create mode 100644 model/common/tests/mpi_tests/test_parallel_setup.py create mode 100644 model/common/tests/test_grid_manager.py create mode 100644 model/common/tests/test_icon_grid.py create mode 100644 model/common/tests/test_interpolation_fields.py rename model/{atmosphere/dycore => common}/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py (85%) create mode 100644 model/common/tests/test_vertical.py create mode 100644 model/driver/.bumpversion.cfg create mode 100644 model/driver/.flake8 create mode 100644 model/driver/.pre-commit-config.yaml create mode 100644 model/driver/README.md create mode 100644 model/driver/pyproject.toml create mode 100644 model/driver/requirements-dev.txt create mode 100644 model/driver/requirements.txt create mode 100644 model/driver/src/icon4py/model/driver/__init__.py create mode 100644 model/driver/src/icon4py/model/driver/dycore_driver.py create mode 100644 model/driver/src/icon4py/model/driver/icon_configuration.py create mode 100644 model/driver/src/icon4py/model/driver/io_utils.py create mode 100644 model/driver/tests/conftest.py create mode 100644 model/driver/tests/test_io_utils.py create mode 100644 tools/src/icon4pytools/py2f/__init__.py create mode 100644 tools/src/icon4pytools/py2f/cffi_utils.py create mode 100644 tools/src/icon4pytools/py2f/codegen.py create mode 100644 tools/src/icon4pytools/py2f/parsing.py create mode 100644 tools/src/icon4pytools/py2f/py2fgen.py create mode 100644 tools/src/icon4pytools/py2f/typing_utils.py create mode 100644 tools/src/icon4pytools/py2f/wrappers/__init__.py create mode 100644 tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py create mode 100644 tools/tests/py2f/test_cffi_utils.py create mode 100644 tools/tests/py2f/test_code_generation.py create mode 100644 tools/tests/py2f/test_parsing_wrapper.py create mode 100644 tools/tests/py2f/test_py2fgen.py diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 0b7bfdb7eb..ab6c89cae3 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -21,10 +21,14 @@ on: types: [submitted] jobs: - pre-commit-icon4py-model-common: + pre-commit-icon4py-model: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: @@ -36,8 +40,8 @@ jobs: **/base-requirements-dev.txt **/requirements.txt **/requirements-dev.txt - - name: Install icon4py-common - working-directory: model/common + - name: Install icon4py-model packages + working-directory: model run: | python -m pip install --upgrade pip setuptools wheel python -m pip install -r ./requirements-dev.txt @@ -45,51 +49,16 @@ jobs: - name: Run checks icon4py-model-common run: | pre-commit run --config model/common/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-atmosphere-dycore: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-atmosphere-dycore - working-directory: model/atmosphere/dycore + - name: Run checks icon4py-model-driver run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt + pre-commit run --config model/driver/.pre-commit-config.yaml --all-files - name: Run checks icon4py-model-atmosphere-dycore run: | pre-commit run --config model/atmosphere/dycore/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-atmosphere-advection: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-atmosphere-advection - working-directory: model/atmosphere/advection + - name: Run checks icon4py-model-atmosphere-diffusion run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt + pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files + - name: Run checks icon4py-model-atmosphere-advection run: | pre-commit run --config model/atmosphere/advection/.pre-commit-config.yaml --all-files diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index 0f550ce84d..af1138fed7 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -31,6 +31,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install boost + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: diff --git a/.gitignore b/.gitignore index 56ff4f8b6c..e7620d5cf3 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,8 @@ _local _external_src _reports tmp +testdata +simple_mesh*.nc ### GT4Py #### .gt_cache/ diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 3c2f34e7af..991b3d4732 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,7 @@ # VCS -e git+https://github.com/GridTools/gt4py.git@main#egg=gt4py +git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python +git+https://github.com/ghex-org/GHEX.git#subdirectory=bindings/python # PyPI bump2version>=1.0.1 @@ -13,6 +15,7 @@ flake8-eradicate>=1.3.0 flake8-mutable>=1.2.0 isort~=5.10 mypy>=0.942 +typing-extensions==4.5.0 pre-commit~=2.15 pytest>=6.1 pytest-benchmark>=4.0.0 @@ -20,6 +23,8 @@ pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 pytest-xdist[psutil]>=2.2 +pytest-mpi>=0.6 setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 +wget>=3.2 diff --git a/base-requirements.txt b/base-requirements.txt index 4e204ea2f5..0720ba2583 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,2 +1,3 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main +pyghex @ git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python diff --git a/model/README.md b/model/README.md index 3f120fbe13..9654b9fadb 100644 --- a/model/README.md +++ b/model/README.md @@ -5,8 +5,10 @@ This folder contains Python implementations for multiple ICON components. It includes the following packages: - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model +- `atmosphere/diffusion`: Contains the implementation of diffusion in the ICON model - `atmosphere/advection`: Contains implementations of the advection component of the ICON model - `common`: Contains shared functionality that is required by multiple components. +- `driver`: Contains the driving code for the model ## Installation Instructions @@ -24,3 +26,47 @@ pip install -r requirements-dev.txt ``` **Note**: For more information specific to each component, please refer to the README in their respective subfolders. + +### Testing + +See the repository [README](../README.md) for general information. + +#### Data dependent tests + +Some test depend on serialized data generated by a full model run. +Those test are marked with `pytest.mark.datatest` and are only run when the `--datatest` +option is specified. Note that due to `pytests` configuration discovery +you need to specify the base a package directory i.e. one that contains a `pyproject.toml` for the +commandline option to work. + +```bash + +pytest -v --datatest model/common +pytest -v --datatest model/atmosphere/diffusion +``` + +#### Testing parallel code + +Tests for parallel codes using MPI are collected in specific subpackages of the model components test folders (e.g. `diffusion_tests/mpi_tests`). All parallel tests are marked with `@pytest.mark.mpi` and are skipped if the `--with-mpi` is not passed option is not passed to `pytest` In order to run them, you need a MPI installation on your system: On Debian-Linux do + +```bash +sudo apt-get install libopenmpi-dev +``` + +or + +```bash +sudo apt-get install mpich +``` + +on MacOs + +```bash +brew install mpich +``` + +Then you can run the tests with + +```bash +mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests +``` diff --git a/model/atmosphere/diffusion/.bumpversion.cfg b/model/atmosphere/diffusion/.bumpversion.cfg new file mode 100644 index 0000000000..ef8bea214e --- /dev/null +++ b/model/atmosphere/diffusion/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/atmosphere/diffusion/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/atmosphere/diffusion/.flake8 b/model/atmosphere/diffusion/.flake8 new file mode 100644 index 0000000000..31cecff5ab --- /dev/null +++ b/model/atmosphere/diffusion/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/atmosphere/diffusion/.pre-commit-config.yaml b/model/atmosphere/diffusion/.pre-commit-config.yaml new file mode 100644 index 0000000000..1d6ce2a032 --- /dev/null +++ b/model/atmosphere/diffusion/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/atmosphere/diffusion/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/atmosphere/diffusion/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/atmosphere/diffusion/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/atmosphere/diffusion/.flake8, model/atmosphere/diffusion/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/atmosphere/diffusion/README.md b/model/atmosphere/diffusion/README.md new file mode 100644 index 0000000000..51a4153193 --- /dev/null +++ b/model/atmosphere/diffusion/README.md @@ -0,0 +1,9 @@ +# icon4py-atmosphere-diffusion + +## Description + +Python port of ICON diffusion module. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/atmosphere/diffusion/diffusion_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/atmosphere/diffusion/diffusion_tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py new file mode 100644 index 0000000000..d429204442 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/conftest.py @@ -0,0 +1,94 @@ +# 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 + +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 # import fixtures from test_utils package + backend, + damping_height, + data_provider, + datapath, + decomposition_info, + download_ser_data, + grid_savepoint, + icon_grid, + interpolation_savepoint, + linit, + mesh, + metrics_savepoint, + ndyn_substeps, + processor_props, + ranked_data_path, + step_date_exit, + step_date_init, +) + + +@pytest.fixture +def r04b09_diffusion_config( + ndyn_substeps, # noqa: F811 # imported `ndyn_substeps` fxiture +) -> DiffusionConfig: + """ + Create DiffusionConfig matching MCH_CH_r04b09_dsl. + + Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + from the default. + """ + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + n_substeps=ndyn_substeps, + ) + + +@pytest.fixture +def diffusion_savepoint_init( + data_provider, # noqa: F811 # imported fixtures data_provider + linit, # noqa: F811 # imported fixtures linit + step_date_init, # noqa: F811 # imported fixtures data_provider +): + """ + Load data from ICON savepoint at start of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' + fixture, passing 'step_date_init=' + + linit flag can be set by overriding the 'linit' fixture + """ + return data_provider.from_savepoint_diffusion_init(linit=linit, date=step_date_init) + + +@pytest.fixture +def diffusion_savepoint_exit( + data_provider, # noqa: F811 # imported fixtures data_provider` + linit, # noqa: F811 # imported fixtures linit` + step_date_exit, # noqa: F811 # imported fixtures step_date_exit` +): + """ + Load data from ICON savepoint at exist of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_data' + fixture, passing 'step_data=' + """ + sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) + return sp diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py new file mode 100644 index 0000000000..8a6602799d --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -0,0 +1,110 @@ +# 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 + + +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size + +from ..utils import verify_diffusion_fields + + +@pytest.mark.mpi +@pytest.mark.parametrize("ndyn_substeps", [2]) +@pytest.mark.parametrize("linit", [True, False]) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_parallel_diffusion( + r04b09_diffusion_config, + step_date_init, + linit, + ndyn_substeps, + processor_props, + decomposition_info, + icon_grid, + diffusion_savepoint_init, + diffusion_savepoint_exit, + grid_savepoint, + metrics_savepoint, + interpolation_savepoint, + damping_height, +): + check_comm_size(processor_props) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: decomposition info : klevels = {decomposition_info.klevels}, " + f"local cells = {decomposition_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " + f"local vertices = {decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + ) + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + + diffusion_params = DiffusionParams(r04b09_diffusion_config) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + exchange = create_exchange(processor_props, decomposition_info) + + diffusion = Diffusion(exchange) + + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=diffusion_params, + vertical_params=VerticalModelParams(grid_savepoint.vct_a(), damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + if linit: + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + else: + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + + verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" + ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py similarity index 95% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index d9d058adf3..02e618fcd8 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( apply_nabla2_and_nabla4_global_to_vn, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py similarity index 62% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py index 4b4a9ceccf..29d3fa3916 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -14,7 +14,10 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( + apply_nabla2_and_nabla4_global_to_vn, +) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( apply_nabla2_and_nabla4_to_vn, ) from icon4py.model.common.dimension import EdgeDim, KDim @@ -67,3 +70,34 @@ def reference( - diff_multfac_vn * z_nabla4_e2 * area_edge ) return dict(vn=vn) + + +class TestApplyNabla2AndNabla4ToVnGlobalMode(StencilTest): + PROGRAM = apply_nabla2_and_nabla4_global_to_vn + + OUTPUTS = ("vn",) + + @pytest.fixture + def input_data(self, mesh): + area_edge = random_field(mesh, EdgeDim) + kh_smag_e = random_field(mesh, EdgeDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) + diff_multfac_vn = random_field(mesh, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + return dict( + area_edge=area_edge, + kh_smag_e=kh_smag_e, + z_nabla2_e=z_nabla2_e, + z_nabla4_e2=z_nabla4_e2, + diff_multfac_vn=diff_multfac_vn, + vn=vn, + ) + + @staticmethod + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py similarity index 94% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index 7928824322..ff44a1f676 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( apply_nabla2_to_vn_in_lateral_boundary, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py similarity index 80% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py index c3f4725214..147fd5c599 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -30,7 +31,8 @@ def reference( z_nabla2_c: np.array, geofac_n2s: np.array, w: np.array, - diff_multfac_w, + diff_multfac_w: float, + **kwargs, ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) @@ -43,12 +45,14 @@ def input_data(self, mesh): z_nabla2_c = random_field(mesh, CellDim, KDim) geofac_n2s = random_field(mesh, CellDim, C2E2CODim) w = random_field(mesh, CellDim, KDim) - diff_multfac_w = 5.0 - return dict( area=area, z_nabla2_c=z_nabla2_c, geofac_n2s=geofac_n2s, w=w, - diff_multfac_w=diff_multfac_w, + diff_multfac_w=5.0, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py similarity index 83% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 98193e0585..6608b90089 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, ) from icon4py.model.common.dimension import CellDim, KDim @@ -37,6 +38,10 @@ def input_data(self, mesh): diff_multfac_n2w=diff_multfac_n2w, cell_area=cell_area, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) @staticmethod @@ -46,6 +51,7 @@ def reference( diff_multfac_n2w: np.array, cell_area: np.array, z_nabla2_c: np.array, + **kwargs, ) -> np.array: cell_area = np.expand_dims(cell_area, axis=-1) w = w + diff_multfac_n2w * cell_area * z_nabla2_c diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py similarity index 94% rename from model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py index c628918570..c70fc8fa27 100644 --- a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( calculate_diagnostics_for_turbulence, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py similarity index 85% rename from model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py index f846a7a788..c7a7016f55 100644 --- a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( calculate_horizontal_gradients_for_turbulence, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim @@ -50,4 +51,8 @@ def input_data(self, mesh): geofac_grg_y=geofac_grg_y, dwdx=dwdx, dwdy=dwdy, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 98% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 184dba0dca..b230948b51 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_and_smag_coefficients_for_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py similarity index 81% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py index 45e4327a9c..310a91c21b 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py @@ -13,8 +13,11 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import calculate_nabla2_for_w +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( + calculate_nabla2_for_w, +) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -39,4 +42,8 @@ def input_data(self, mesh): w=w, geofac_n2s=geofac_n2s, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py similarity index 94% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py index 7772dd5dbd..a82916ce3b 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py @@ -14,7 +14,9 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import calculate_nabla2_for_z +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + calculate_nabla2_for_z, +) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py similarity index 93% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py index 5b1f998750..0d6bd52515 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py @@ -14,7 +14,9 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import calculate_nabla2_of_theta +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + calculate_nabla2_of_theta, +) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( StencilTest, diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py similarity index 97% rename from model/atmosphere/dycore/tests/test_calculate_nabla4.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py index 07d097bf16..de26d3cd6e 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla4 import calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.model.common.test_utils.helpers import ( StencilTest, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py new file mode 100644 index 0000000000..383d10f342 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -0,0 +1,348 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionInitSavepoint + +from .utils import ( + diff_multfac_vn_numpy, + enhanced_smagorinski_factor_numpy, + smag_limit_numpy, + verify_diffusion_fields, +) + + +def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.hdiff_efdt_ratio = 1.0 + config.hdiff_w_efdt_ratio = 2.0 + + params = DiffusionParams(config) + + assert params.K2 == pytest.approx(0.125, abs=1e-12) + assert params.K4 == pytest.approx(0.125 / 8.0, abs=1e-12) + assert params.K6 == pytest.approx(0.125 / 64.0, abs=1e-12) + assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) + + +def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.hdiff_efdt_ratio = 0.0 + config.hdiff_w_efdt_ratio = 0.0 + + params = DiffusionParams(config) + + assert params.K2 == 0.0 + assert params.K4 == 0.0 + assert params.K6 == 0.0 + assert params.K4W == 0.0 + + +def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.smagorinski_scaling_factor = 0.15 + config.diffusion_type = 4 + + params = DiffusionParams(config) + assert len(params.smagorinski_factor) == 1 + assert params.smagorinski_factor[0] == pytest.approx(0.15, abs=1e-16) + assert params.smagorinski_height is None + + +def test_smagorinski_heights_diffusion_type_5_are_consistent( + r04b09_diffusion_config, +): + config = r04b09_diffusion_config + config.smagorinski_scaling_factor = 0.15 + config.diffusion_type = 5 + + params = DiffusionParams(config) + assert len(params.smagorinski_height) == 4 + assert min(params.smagorinski_height) == params.smagorinski_height[0] + assert max(params.smagorinski_height) == params.smagorinski_height[-1] + assert params.smagorinski_height[0] < params.smagorinski_height[1] + assert params.smagorinski_height[1] < params.smagorinski_height[3] + assert params.smagorinski_height[2] != params.smagorinski_height[1] + assert params.smagorinski_height[2] != params.smagorinski_height[3] + + +def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): + params = DiffusionParams(r04b09_diffusion_config) + assert len(params.smagorinski_factor) == len(params.smagorinski_height) + assert len(params.smagorinski_factor) == 4 + assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) + + +@pytest.mark.datatest +def test_diffusion_init( + diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + step_date_init, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + + meta = diffusion_savepoint_init.get_metadata("linit", "date") + + assert meta["linit"] is False + assert meta["date"] == step_date_init + + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + assert diffusion.diff_multfac_w == min( + 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float + ) + + assert np.allclose(0.0, np.asarray(diffusion.v_vert)) + assert np.allclose(0.0, np.asarray(diffusion.u_vert)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) + + shape_k = (icon_grid.n_lev(),) + expected_smag_limit = smag_limit_numpy( + diff_multfac_vn_numpy, + shape_k, + additional_parameters.K4, + config.substep_as_float, + ) + + assert diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float + assert np.allclose(expected_smag_limit, diffusion.smag_limit) + + expected_diff_multfac_vn = diff_multfac_vn_numpy( + shape_k, additional_parameters.K4, config.substep_as_float + ) + assert np.allclose(expected_diff_multfac_vn, diffusion.diff_multfac_vn) + expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( + additional_parameters.smagorinski_factor, + additional_parameters.smagorinski_height, + grid_savepoint.vct_a(), + ) + assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) + + +def _verify_init_values_against_savepoint( + savepoint: IconDiffusionInitSavepoint, diffusion: Diffusion +): + dtime = savepoint.get_metadata("dtime")["dtime"] + + assert savepoint.nudgezone_diff() == diffusion.nudgezone_diff + assert savepoint.bdy_diff() == diffusion.bdy_diff + assert savepoint.fac_bdydiff_v() == diffusion.fac_bdydiff_v + assert savepoint.smag_offset() == diffusion.smag_offset + assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w + + # this is done in diffusion.run(...) because it depends on the dtime + scale_k(diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={}) + assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) + + assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) + assert np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) + + +@pytest.mark.datatest +def test_verify_diffusion_init_against_first_regular_savepoint( + diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + r04b09_diffusion_config, + icon_grid, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + vct_a = grid_savepoint.vct_a() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() + + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=VerticalModelParams(vct_a, damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) + + +@pytest.mark.datatest +@pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) +def test_verify_diffusion_init_against_other_regular_savepoint( + r04b09_diffusion_config, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + diffusion_savepoint_init, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() + + diffusion = Diffusion() + diffusion.init( + icon_grid, + config, + additional_parameters, + vertical_params, + metric_state, + interpolation_state, + edge_params, + cell_params, + ) + + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "step_date_init, step_date_exit", + [ + ("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), + ], +) +def test_run_diffusion_single_step( + diffusion_savepoint_init, + diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, +): + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + vct_a = grid_savepoint.vct_a() + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_init) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit) + + +@pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) +def test_run_diffusion_initial_step( + diffusion_savepoint_init, + diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, +): + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + vct_a = grid_savepoint.vct_a() + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + + verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py new file mode 100644 index 0000000000..243ab01e2c --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py @@ -0,0 +1,28 @@ +# 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 + +import numpy as np +import pytest + + +@pytest.mark.datatest +def test_verify_geofac_n2s_field_manipulation(interpolation_savepoint, icon_grid): + geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) + int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + geofac_c = np.asarray(int_state.geofac_n2s_c) + geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) + assert np.count_nonzero(geofac_nbh) > 0 + cec_table = icon_grid.get_c2cec_connectivity().table + assert np.allclose(geofac_c, geofac_n2s[:, 0]) + assert geofac_nbh[cec_table].shape == geofac_n2s[:, 1:].shape + assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py new file mode 100644 index 0000000000..0ba7ce353f --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py @@ -0,0 +1,147 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + _en_smag_fac_for_zero_nshift, + _setup_runtime_diff_multfac_vn, + _setup_smag_limit, + scale_k, + set_zero_v_k, + setup_fields_for_initial_step, +) +from icon4py.model.common.dimension import KDim, VertexDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + +from .utils import diff_multfac_vn_numpy, enhanced_smagorinski_factor_numpy, smag_limit_numpy + + +def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): + return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) + + +def test_scale_k(): + mesh = SimpleMesh() + field = random_field(mesh, KDim) + scaled_field = zero_field(mesh, KDim) + factor = 2.0 + scale_k(field, factor, scaled_field, offset_provider={}) + assert np.allclose(factor * np.asarray(field), scaled_field) + + +def test_diff_multfac_vn_and_smag_limit_for_initial_step(): + mesh = SimpleMesh() + diff_multfac_vn_init = zero_field(mesh, KDim) + smag_limit_init = zero_field(mesh, KDim) + k4 = 1.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn_init).shape + + expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio) + expected_smag_limit_init = smag_limit_numpy( + initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio + ) + + setup_fields_for_initial_step( + k4, efdt_ratio, diff_multfac_vn_init, smag_limit_init, offset_provider={} + ) + + assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) + assert np.allclose(expected_smag_limit_init, smag_limit_init) + + +def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 1.0 + substeps = 5.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn).shape + + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + + _setup_runtime_diff_multfac_vn(k4, efdt_ratio, out=diff_multfac_vn, offset_provider={}) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 0.003 + substeps = 1.0 + + shape = np.asarray(diff_multfac_vn).shape + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + _setup_runtime_diff_multfac_vn(k4, substeps, out=diff_multfac_vn, offset_provider={}) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def test_init_enh_smag_fac(): + mesh = SimpleMesh() + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0, extend={KDim: 1}) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + + enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) + + _en_smag_fac_for_zero_nshift(a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim}) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) + + +def test_set_zero_vertex_k(): + mesh = SimpleMesh() + f = random_field(mesh, VertexDim, KDim) + set_zero_v_k(f, offset_provider={}) + assert np.allclose(0.0, f) + + +@pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) +def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, linit +): + savepoint = diffusion_savepoint_init + config = r04b09_diffusion_config + + params = DiffusionParams(config) + expected_diff_multfac_vn = savepoint.diff_multfac_vn() + expected_smag_limit = savepoint.smag_limit() + exptected_smag_offset = savepoint.smag_offset() + + diff_multfac_vn = zero_field(icon_grid, KDim) + smag_limit = zero_field(icon_grid, KDim) + setup_fields_for_initial_step( + params.K4, + config.hdiff_efdt_ratio, + diff_multfac_vn, + smag_limit, + offset_provider={}, + ) + assert np.allclose(expected_smag_limit, smag_limit) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert exptected_smag_offset == 0.0 diff --git a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 92% rename from model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 50e92d4f58..5541422269 100644 --- a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( enhance_diffusion_coefficient_for_grid_point_cold_pools, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 94% rename from model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index 61966f1ad8..84d8ed5160 100644 --- a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py similarity index 74% rename from model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py index 43cd3f4aa2..6b044d2fbc 100644 --- a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -14,11 +14,16 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + zero_field, +) class TestTemporaryFieldsForTurbulenceDiagnostics(StencilTest): @@ -36,12 +41,11 @@ def reference( **kwargs, ) -> dict: geofac_div = np.expand_dims(geofac_div, axis=-1) - vn_geofac = vn[mesh.c2e] * geofac_div + vn_geofac = vn[mesh.c2e] * geofac_div[mesh.get_c2ce_offset_provider().table] div = np.sum(vn_geofac, axis=1) - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) diff_multfac_smag = np.expand_dims(diff_multfac_smag, axis=0) - mul = kh_smag_ec[mesh.c2e] * e_bln_c_s + mul = kh_smag_ec[mesh.c2e] * e_bln_c_s[mesh.get_c2ce_offset_provider().table] summed = np.sum(mul, axis=1) kh_c = summed / diff_multfac_smag @@ -50,9 +54,9 @@ def reference( @pytest.fixture def input_data(self, mesh): vn = random_field(mesh, EdgeDim, KDim) - geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) kh_smag_ec = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) + e_bln_c_s = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) diff_multfac_smag = random_field(mesh, KDim) kh_c = zero_field(mesh, CellDim, KDim) diff --git a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 93% rename from model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index b0114cc26a..72baed3943 100644 --- a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim @@ -111,6 +111,10 @@ def test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points(): vcoef_new, theta_v, z_temp, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), diff --git a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py similarity index 84% rename from model/atmosphere/dycore/tests/test_update_theta_and_exner.py rename to model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py index 7292543373..871b32a894 100644 --- a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py @@ -13,8 +13,11 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.update_theta_and_exner import update_theta_and_exner +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + update_theta_and_exner, +) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -53,4 +56,8 @@ def input_data(self, mesh): theta_v=theta_v, exner=exner, rd_o_cvd=rd_o_cvd, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py new file mode 100644 index 0000000000..b2fbf987d5 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -0,0 +1,77 @@ +# 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 + +import numpy as np + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionExitSavepoint + + +def verify_diffusion_fields( + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0 / 128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + + +def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml new file mode 100644 index 0000000000..d6ab9aa52f --- /dev/null +++ b/model/atmosphere/diffusion/pyproject.toml @@ -0,0 +1,120 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] + +[project] +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "icon4py-common>=0.0.5", + "mpi4py<=3.1.4", + "pyghex>=0.3.0" +] +description = "ICON diffusion." +dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-atmosphere-diffusion" +readme = "README.md" +requires-python = ">=3.10" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'diffusion_tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +branch = true +parallel = true +source_pkgs = ['diffusion'] + +[tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +multi_line_output = 3 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +use_parentheses = true + +[tool.mypy] +disallow_incomplete_defs = true +disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] +ignore_missing_imports = false +implicit_reexport = true +install_types = true +non_interactive = true +show_column_numbers = true +show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = ['tests', 'diffusion_tests'] + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.atmosphere.diffusion.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.atmosphere.diffusion' = ['py.typed'] diff --git a/model/atmosphere/diffusion/requirements-dev.txt b/model/atmosphere/diffusion/requirements-dev.txt new file mode 100644 index 0000000000..32d385e08d --- /dev/null +++ b/model/atmosphere/diffusion/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements-dev.txt +-e ../../common +-e . diff --git a/model/atmosphere/diffusion/requirements.txt b/model/atmosphere/diffusion/requirements.txt new file mode 100644 index 0000000000..79787ec137 --- /dev/null +++ b/model/atmosphere/diffusion/requirements.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements.txt +../../common +. diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py new file mode 100644 index 0000000000..49c96a94c7 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py @@ -0,0 +1,34 @@ +# 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 typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py new file mode 100644 index 0000000000..d857a4f275 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -0,0 +1,850 @@ +# 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 +import functools +import logging +import math +import sys +from collections import namedtuple +from dataclasses import InitVar, dataclass, field +from enum import Enum +from typing import Final, Optional + +import numpy as np +from gt4py.next.common import Dimension +from gt4py.next.ffront.fbuiltins import Field, int32 +from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next.program_processors.runners.gtfn_cpu import ( + run_gtfn, + run_gtfn_cached, + run_gtfn_imperative, +) + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + copy_field, + init_diffusion_local_fields_for_regular_timestep, + init_nabla2_factor_in_upper_damping_zone, + scale_k, + setup_fields_for_initial_step, + zero_field, +) +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostic_quantities_for_turbulence import ( + calculate_diagnostic_quantities_for_turbulence, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( + calculate_nabla2_and_smag_coefficients_for_vn, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_theta import ( + calculate_nabla2_for_theta, +) +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points, +) +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + update_theta_and_exner, +) +from icon4py.model.common.constants import ( + CPD, + DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + GAS_CONSTANT_DRY_AIR, +) +from icon4py.model.common.decomposition.decomposed import ExchangeRuntime, SingleNode +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex +from icon4py.model.common.grid.icon_grid import IconGrid +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( + mo_intp_rbf_rbf_vec_interpol_vertex, +) + + +# flake8: noqa +log = logging.getLogger(__name__) + +VectorTuple = namedtuple("VectorTuple", "x y") + +cached_backend = run_gtfn_cached +compiled_backend = run_gtfn +imperative_backend = run_gtfn_imperative +backend = run_gtfn_cached # + + +class DiffusionType(int, Enum): + """ + Order of nabla operator for diffusion. + Note: Called `hdiff_order` in `mo_diffusion_nml.f90`. + Note: We currently only support type 5. + """ + + NO_DIFFUSION = -1 #: no diffusion + LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels + SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion + LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels + SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion + + +class DiffusionConfig: + """ + Contains necessary parameter to configure a diffusion run. + + Encapsulates namelist parameters and derived parameters. + Values should be read from configuration. + Default values are taken from the defaults in the corresponding ICON Fortran namelist files. + """ + + # TODO(Magdalena): to be read from config + # TODO(Magdalena): handle dependencies on other namelists (see below...) + + def __init__( + self, + diffusion_type: DiffusionType = DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + hdiff_temp=True, + type_vn_diffu: int = 1, + smag_3d: bool = False, + type_t_diffu: int = 2, + hdiff_efdt_ratio: float = 36.0, + hdiff_w_efdt_ratio: float = 15.0, + smagorinski_scaling_factor: float = 0.015, + n_substeps: int = 5, + zdiffu_t: bool = True, + hdiff_rcf: bool = True, + velocity_boundary_diffusion_denom: float = 200.0, + temperature_boundary_diffusion_denom: float = 135.0, + max_nudging_coeff: float = 0.02, + nudging_decay_rate: float = 2.0, + ): + """Set the diffusion configuration parameters with the ICON default values.""" + # parameters from namelist diffusion_nml + + self.diffusion_type: int = diffusion_type + + #: If True, apply diffusion on the vertical wind field + #: Called `lhdiff_w` in mo_diffusion_nml.f90 + self.apply_to_vertical_wind: bool = hdiff_w + + #: True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90 + #: Called `lhdiff_vn` in mo_diffusion_nml.f90 + self.apply_to_horizontal_wind = hdiff_vn + + #: If True, apply horizontal diffusion to temperature field + #: Called `lhdiff_temp` in mo_diffusion_nml.f90 + self.apply_to_temperature: bool = hdiff_temp + + #: If True, compute 3D Smagorinsky diffusion coefficient + #: Called `lsmag_3d` in mo_diffusion_nml.f90 + self.compute_3d_smag_coeff: bool = smag_3d + + #: Options for discretizing the Smagorinsky momentum diffusion + #: Called `itype_vn_diffu` in mo_diffusion_nml.f90 + self.type_vn_diffu: int = type_vn_diffu + + #: Options for discretizing the Smagorinsky temperature diffusion + #: Called `itype_t_diffu` inmo_diffusion_nml.f90 + self.type_t_diffu = type_t_diffu + + #: Ratio of e-folding time to (2*)time step + #: Called `hdiff_efdt_ratio` inmo_diffusion_nml.f90 + self.hdiff_efdt_ratio: float = hdiff_efdt_ratio + + #: Ratio of e-folding time to time step for w diffusion (NH only) + #: Called `hdiff_w_efdt_ratio` inmo_diffusion_nml.f90. + self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio + + #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below + #: Called `hdiff_smag_fac` inmo_diffusion_nml.f90 + self.smagorinski_scaling_factor: float = smagorinski_scaling_factor + + #: If True, apply truly horizontal temperature diffusion over steep slopes + #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90 + self.apply_zdiffusion_t: bool = zdiffu_t + + # from other namelists: + # from parent namelist mo_nonhydrostatic_nml + + #: Number of dynamics substeps per fast-physics step + #: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90 + self.ndyn_substeps: int = n_substeps + + #: If True, compute horizontal diffusion only at the large time step + #: Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90 + self.lhdiff_rcf: bool = hdiff_rcf + + # namelist mo_gridref_nml.f90 + + #: Denominator for temperature boundary diffusion + #: Called 'denom_diffu_t' in mo_gridref_nml.f90 + self.temperature_boundary_diffusion_denominator: float = ( + temperature_boundary_diffusion_denom + ) + + #: Denominator for velocity boundary diffusion + #: Called 'denom_diffu_v' in mo_gridref_nml.f90 + self.velocity_boundary_diffusion_denominator: float = velocity_boundary_diffusion_denom + + # parameters from namelist: mo_interpol_nml.f90 + + #: Parameter describing the lateral boundary nudging in limited area mode. + #: + #: Maximal value of the nudging coefficients used cell row bordering the boundary interpolation zone, + #: from there nudging coefficients decay exponentially with `nudge_efold_width` in units of cell rows. + #: Called `nudge_max_coeff` in mo_interpol_nml.f90 + self.nudge_max_coeff: float = max_nudging_coeff + + #: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients + #: Called `nudge_efold_width` in mo_interpol_nml.f90 + self.nudge_efold_width: float = nudging_decay_rate + + self._validate() + + def _validate(self): + """Apply consistency checks and validation on configuration parameters.""" + if self.diffusion_type != 5: + raise NotImplementedError( + "Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background " + "diffusion` is implemented" + ) + + if self.diffusion_type < 0: + self.apply_to_temperature = False + self.apply_to_horizontal_wind = False + self.apply_to_vertical_wind = False + else: + self.apply_to_temperature = True + self.apply_to_horizontal_wind = True + + if not self.apply_zdiffusion_t: + raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)") + + @functools.cached_property + def substep_as_float(self): + return float(self.ndyn_substeps) + + +@dataclass(frozen=True) +class DiffusionParams: + """Calculates derived quantities depending on the diffusion config.""" + + config: InitVar[DiffusionConfig] + K2: Final[float] = field(init=False) + K4: Final[float] = field(init=False) + K6: Final[float] = field(init=False) + K4W: Final[float] = field(init=False) + smagorinski_factor: Final[float] = field(init=False) + smagorinski_height: Final[float] = field(init=False) + scaled_nudge_max_coeff: Final[float] = field(init=False) + + def __post_init__(self, config): + object.__setattr__( + self, + "K2", + (1.0 / (config.hdiff_efdt_ratio * 8.0) if config.hdiff_efdt_ratio > 0.0 else 0.0), + ) + object.__setattr__(self, "K4", self.K2 / 8.0) + object.__setattr__(self, "K6", self.K2 / 64.0) + object.__setattr__( + self, + "K4W", + (1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0), + ) + + ( + smagorinski_factor, + smagorinski_height, + ) = self._determine_smagorinski_factor(config) + object.__setattr__(self, "smagorinski_factor", smagorinski_factor) + object.__setattr__(self, "smagorinski_height", smagorinski_height) + # see mo_interpol_nml.f90: + object.__setattr__( + self, + "scaled_nudge_max_coeff", + config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + ) + + def _determine_smagorinski_factor(self, config: DiffusionConfig): + """Enhanced Smagorinsky diffusion factor. + + Smagorinsky diffusion factor is defined as a profile in height + above sea level with 4 height sections. + + It is calculated/used only in the case of diffusion_type 3 or 5 + """ + match config.diffusion_type: + case 5: + ( + smagorinski_factor, + smagorinski_height, + ) = diffusion_type_5_smagorinski_factor(config) + case 4: + # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only + # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 + smagorinski_factor = ( + config.smagorinski_scaling_factor + if config.smagorinski_scaling_factor + else 0.15, + ) + smagorinski_height = None + case _: + raise NotImplementedError("Only implemented for diffusion type 4 and 5") + smagorinski_factor = None + smagorinski_height = None + pass + return smagorinski_factor, smagorinski_height + + +def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): + """ + Initialize Smagorinski factors used in diffusion type 5. + + The calculation and magic numbers are taken from mo_diffusion_nml.f90 + """ + magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) + magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) + magic_z2 = 1600.0 + 50000.0 + magic_sqrt + factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) + heights = (32500.0, magic_z2, 50000.0, 90000.0) + return factor, heights + + +class Diffusion: + """Class that configures diffusion and does one diffusion step.""" + + def __init__(self, exchange: ExchangeRuntime = SingleNode()): + self._exchange = exchange + self._initialized = False + self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) + self.thresh_tdiff: float = ( + -5.0 + ) #: threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + self.grid: Optional[IconGrid] = None + self.config: Optional[DiffusionConfig] = None + self.params: Optional[DiffusionParams] = None + self.vertical_params: Optional[VerticalModelParams] = None + self.interpolation_state: DiffusionInterpolationState = None + self.metric_state: DiffusionMetricState = None + self.diff_multfac_w: Optional[float] = None + self.diff_multfac_n2w: Field[[KDim], float] = None + self.smag_offset: Optional[float] = None + self.fac_bdydiff_v: Optional[float] = None + self.bdy_diff: Optional[float] = None + self.nudgezone_diff: Optional[float] = None + self.edge_params: Optional[EdgeParams] = None + self.cell_params: Optional[CellParams] = None + self._horizontal_start_index_w_diffusion: int32 = 0 + + def init( + self, + grid: IconGrid, + config: DiffusionConfig, + params: DiffusionParams, + vertical_params: VerticalModelParams, + metric_state: DiffusionMetricState, + interpolation_state: DiffusionInterpolationState, + edge_params: EdgeParams, + cell_params: CellParams, + ): + """ + Initialize Diffusion granule with configuration. + + calculates all local fields that are used in diffusion within the time loop. + + Args: + grid: + config: + params: + vertical_params: + metric_state: + interpolation_state: + edge_params: + cell_params: + """ + self.config: DiffusionConfig = config + self.params: DiffusionParams = params + self.grid = grid + self.vertical_params = vertical_params + self.metric_state: DiffusionMetricState = metric_state + self.interpolation_state: DiffusionInterpolationState = interpolation_state + self.edge_params = edge_params + self.cell_params = cell_params + + self._allocate_temporary_fields() + + def _get_start_index_for_w_diffusion() -> int32: + return self.grid.get_start_index( + CellDim, + ( + HorizontalMarkerIndex.nudging(CellDim) + if self.grid.limited_area() + else HorizontalMarkerIndex.interior(CellDim) + ), + ) + + self.nudgezone_diff: float = 0.04 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) + self.bdy_diff: float = 0.015 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) + self.fac_bdydiff_v: float = ( + math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator + if config.lhdiff_rcf + else 1.0 / config.velocity_boundary_diffusion_denominator + ) + + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float + self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float) + + init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( + params.K4, + config.substep_as_float, + *params.smagorinski_factor, + *params.smagorinski_height, + self.vertical_params.physical_heights, + self.diff_multfac_vn, + self.smag_limit, + self.enh_smag_fac, + offset_provider={"Koff": KDim}, + ) + + self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( + k_size=self.grid.n_lev(), + nshift=0, + physical_heights=np.asarray(self.vertical_params.physical_heights), + nrdmax=self.vertical_params.index_of_damping_layer, + ) + self._horizontal_start_index_w_diffusion = _get_start_index_for_w_diffusion() + self._initialized = True + + @property + def initialized(self): + return self._initialized + + def _allocate_temporary_fields(self): + def _allocate(*dims: Dimension): + return zero_field(self.grid, *dims) + + def _index_field(dim: Dimension, size=None): + size = size if size else self.grid.size[dim] + return np_as_located_field(dim)(np.arange(size, dtype=int32)) + + self.diff_multfac_vn = _allocate(KDim) + + self.smag_limit = _allocate(KDim) + self.enh_smag_fac = _allocate(KDim) + self.u_vert = _allocate(VertexDim, KDim) + self.v_vert = _allocate(VertexDim, KDim) + self.kh_smag_e = _allocate(EdgeDim, KDim) + self.kh_smag_ec = _allocate(EdgeDim, KDim) + self.z_nabla2_e = _allocate(EdgeDim, KDim) + self.z_temp = _allocate(CellDim, KDim) + self.diff_multfac_smag = _allocate(KDim) + self.z_nabla4_e2 = _allocate(EdgeDim, KDim) + # TODO(Magdalena): this is KHalfDim + self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) + self.horizontal_cell_index = _index_field(CellDim) + self.horizontal_edge_index = _index_field(EdgeDim) + self.w_tmp = np_as_located_field(CellDim, KDim)( + np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float) + ) + + def initial_run( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + ): + """ + Calculate initial diffusion step. + + In ICON at the start of the simulation diffusion is run with a parameter linit = True: + + 'For real-data runs, perform an extra diffusion call before the first time + step because no other filtering of the interpolated velocity field is done' + + This run uses special values for diff_multfac_vn, smag_limit and smag_offset + + """ + diff_multfac_vn = zero_field(self.grid, KDim) + smag_limit = zero_field(self.grid, KDim) + + setup_fields_for_initial_step.with_backend(backend)( + self.params.K4, + self.config.hdiff_efdt_ratio, + diff_multfac_vn, + smag_limit, + offset_provider={}, + ) + self._do_diffusion_step( + diagnostic_state, + prognostic_state, + dtime, + diff_multfac_vn, + smag_limit, + 0.0, + ) + self._sync_cell_fields(prognostic_state) + + def run( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + ): + """ + Do one diffusion step within regular time loop. + + runs a diffusion step for the parameter linit=False, within regular time loop. + """ + + self._do_diffusion_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + diff_multfac_vn=self.diff_multfac_vn, + smag_limit=self.smag_limit, + smag_offset=self.smag_offset, + ) + if not self.config.lhdiff_rcf: + self._sync_cell_fields(prognostic_state) + + def _sync_cell_fields(self, prognostic_state): + """ + Communicate theta_v, exner and w. + + communication only done in original code if the following condition applies: + IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + """ + log.debug("communication of prognostic cell fields: theta, w, exner - start") + handle_cell_comm = self._exchange.exchange( + CellDim, + prognostic_state.w, + prognostic_state.theta_v, + prognostic_state.exner_pressure, + ) + handle_cell_comm.wait() + log.debug("communication of prognostic cell fields: theta, w, exner - done") + + def _do_diffusion_step( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + ): + """ + Run a diffusion step. + + Args: + diagnostic_state: output argument, data class that contains diagnostic variables + prognostic_state: output argument, data class that contains prognostic variables + dtime: the time step, + diff_multfac_vn: + smag_limit: + smag_offset: + + """ + klevels = self.grid.n_lev() + cell_start_interior = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) + ) + cell_start_nudging = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.nudging(CellDim) + ) + cell_end_local = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) + cell_end_halo = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.halo(CellDim)) + + edge_start_nudging_plus_one = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 + ) + edge_start_nudging = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + ) + edge_start_lb_plus4 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + ) + edge_end_local = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.local(EdgeDim)) + edge_end_local_minus2 = self.grid.get_end_index( + EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 + ) + edge_end_halo = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.halo(EdgeDim)) + + vertex_start_lb_plus1 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 + ) + vertex_end_local = self.grid.get_end_index( + VertexDim, HorizontalMarkerIndex.local(VertexDim) + ) + + # dtime dependent: enh_smag_factor, + scale_k.with_backend(backend)( + self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} + ) + + log.debug("rbf interpolation 1: start") + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + p_e_in=prognostic_state.vn, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + ) + log.debug("rbf interpolation 1: end") + + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert + log.debug("communication rbf extrapolation of vn - start") + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h.wait() + log.debug("communication rbf extrapolation of vn - end") + + log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( + diff_multfac_smag=self.diff_multfac_smag, + tangent_orientation=self.edge_params.tangent_orientation, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_x=self.edge_params.primal_normal_vert[0], + primal_normal_vert_y=self.edge_params.primal_normal_vert[1], + dual_normal_vert_x=self.edge_params.dual_normal_vert[0], + dual_normal_vert_y=self.edge_params.dual_normal_vert[1], + vn=prognostic_state.vn, + smag_limit=smag_limit, + kh_smag_e=self.kh_smag_e, + kh_smag_ec=self.kh_smag_ec, + z_nabla2_e=self.z_nabla2_e, + smag_offset=smag_offset, + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local_minus2, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start") + calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( + kh_smag_ec=self.kh_smag_ec, + vn=prognostic_state.vn, + e_bln_c_s=self.interpolation_state.e_bln_c_s, + geofac_div=self.interpolation_state.geofac_div, + diff_multfac_smag=self.diff_multfac_smag, + wgtfac_c=self.metric_state.wgtfac_c, + div_ic=diagnostic_state.div_ic, + hdef_ic=diagnostic_state.hdef_ic, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=1, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), + "Koff": KDim, + }, + ) + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end") + + # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + + if self.config.type_vn_diffu > 1: + log.debug("communication rbf extrapolation of z_nable2_e - start") + h_z = self._exchange.exchange(EdgeDim, self.z_nabla2_e) + h_z.wait() + log.debug("communication rbf extrapolation of z_nable2_e - end") + + log.debug("2nd rbf interpolation: start") + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + p_e_in=self.z_nabla2_e, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + ) + log.debug("2nd rbf interpolation: end") + + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) + log.debug("communication rbf extrapolation of z_nable2_e - start") + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h.wait() + log.debug("communication rbf extrapolation of z_nable2_e - end") + + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") + apply_diffusion_to_vn.with_backend(backend)( + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], + primal_normal_vert_v2=self.edge_params.primal_normal_vert[1], + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + area_edge=self.edge_params.edge_areas, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + vn=prognostic_state.vn, + horz_idx=self.horizontal_edge_index, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v=self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + limited_area=self.grid.limited_area(), + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") + log.debug("communication of prognistic.vn : start") + handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) + + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + ) + # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify? + copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( + area=self.cell_params.area, + geofac_n2s=self.interpolation_state.geofac_n2s, + geofac_grg_x=self.interpolation_state.geofac_grg_x, + geofac_grg_y=self.interpolation_state.geofac_grg_y, + w_old=self.w_tmp, + w=prognostic_state.w, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=self.vertical_index, + horz_idx=self.horizontal_cell_index, + nrdmax=int32( + self.vertical_params.index_of_damping_layer + 1 + ), # +1 since Fortran includes boundaries + interior_idx=int32(cell_start_interior), + halo_idx=int32(cell_end_local), + horizontal_start=self._horizontal_start_index_w_diffusion, + horizontal_end=cell_end_halo, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity(), + }, + ) + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + ) + + log.debug( + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + ) + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(backend)( + theta_v=prognostic_state.theta_v, + theta_ref_mc=self.metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + horizontal_start=edge_start_nudging, + horizontal_end=edge_end_halo, + vertical_start=(klevels - 2), + vertical_end=klevels, + offset_provider={ + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + }, + ) + log.debug( + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + ) + log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") + calculate_nabla2_for_theta.with_backend(backend)( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths, + theta_v=prognostic_state.theta_v, + geofac_div=self.interpolation_state.geofac_div, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "E2C": self.grid.get_e2c_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), + }, + ) + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") + log.debug( + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + ) + truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)( + mask=self.metric_state.mask_hdiff, + zd_vertoffset=self.metric_state.zd_vertoffset, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + vcoef=self.metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2CEC": self.grid.get_c2cec_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "Koff": KDim, + }, + ) + + log.debug( + "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + ) + log.debug("running stencil 16 (update_theta_and_exner): start") + update_theta_and_exner.with_backend(backend)( + z_temp=self.z_temp, + area=self.cell_params.area, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={}, + ) + log.debug("running stencil 16 (update_theta_and_exner): end") + handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. + log.debug("communication of prognogistic.vn - end") diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py new file mode 100644 index 0000000000..ea423047cb --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -0,0 +1,115 @@ +# 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 +import functools +from dataclasses import dataclass + +import numpy as np +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.common.dimension import ( + C2E2CODim, + CECDim, + CEDim, + CellDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) + + +@dataclass(frozen=True) +class DiffusionDiagnosticState: + """Represents the diagnostic fields needed in diffusion.""" + + # fields for 3D elements in turbdiff + hdef_ic: Field[ + [CellDim, KDim], float + ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] + div_ic: Field[ + [CellDim, KDim], float + ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] + dwdx: Field[ + [CellDim, KDim], float + ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + + dwdy: Field[ + [CellDim, KDim], float + ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) + + +@dataclass(frozen=True) +class DiffusionMetricState: + """Represents the metric state fields needed in diffusion.""" + + theta_ref_mc: Field[[CellDim, KDim], float] + wgtfac_c: Field[ + [CellDim, KDim], float + ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) + mask_hdiff: Field[[CellDim, KDim], bool] + zd_vertoffset: Field[[CECDim, KDim], int32] + zd_diffcoef: Field[[CellDim, KDim], float] + zd_intcoef: Field[[CECDim, KDim], float] + + +@dataclass(frozen=True) +class DiffusionInterpolationState: + """Represents the ICON interpolation state needed in diffusion.""" + + e_bln_c_s: Field[[CEDim], float] # coefficent for bilinear interpolation from edge to cell () + rbf_coeff_1: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + + geofac_div: Field[[CEDim], float] # factor for divergence (nproma,cell_type,nblks_c) + + geofac_n2s: Field[ + [CellDim, C2E2CODim], float + ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) + geofac_grg_x: Field[[CellDim, C2E2CODim], float] + geofac_grg_y: Field[ + [CellDim, C2E2CODim], float + ] # factors for green gauss gradient (nproma,4,nblks_c,2) + nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges + + @functools.cached_property + def geofac_n2s_c(self) -> Field[[CellDim], float]: + return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) + + @functools.cached_property + def geofac_n2s_nbh(self) -> Field[[CECDim], float]: + geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:] + old_shape = geofac_nbh_ar.shape + return np_as_located_field(CECDim)( + geofac_nbh_ar.reshape( + old_shape[0] * old_shape[1], + ) + ) + + +@dataclass +class PrognosticState: + """Class that contains the prognostic state. + + Corresponds to ICON t_nh_prog + """ + + w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py new file mode 100644 index 0000000000..183d146cb1 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -0,0 +1,230 @@ +# 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 typing import Tuple + +import numpy as np +from gt4py.next.common import Dimension, Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, minimum +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim + + +# TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py +def zero_field(mesh, *dims: Dimension, dtype=float): + shapex = tuple(map(lambda x: mesh.size[x], dims)) + return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) + + +@field_operator +def _identity_c_k(field: Field[[CellDim, KDim], float]) -> Field[[CellDim, KDim], float]: + return field + + +@program +def copy_field(old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float]): + _identity_c_k(old_f, out=new_f) + + +@field_operator +def _identity_e_k(field: Field[[EdgeDim, KDim], float]) -> Field[[EdgeDim, KDim], float]: + return field + + +@field_operator +def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: + return field * factor + + +@program +def scale_k(field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float]): + _scale_k(field, factor, out=scaled_field) + + +@field_operator +def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: + return broadcast(0.0, (VertexDim, KDim)) + + +@program +def set_zero_v_k(field: Field[[VertexDim, KDim], float]): + _set_zero_v_k(out=field) + + +@field_operator +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: + return 0.125 - 4.0 * diff_multfac_vn + + +@field_operator +def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float) -> Field[[KDim], float]: + con = 1.0 / 128.0 + dyn = k4 * dyn_substeps / 3.0 + return broadcast(minimum(con, dyn), (KDim,)) + + +@field_operator +def _setup_initial_diff_multfac_vn(k4: float, hdiff_efdt_ratio: float) -> Field[[KDim], float]: + return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) + + +@field_operator +def _setup_fields_for_initial_step( + k4: float, hdiff_efdt_ratio: float +) -> Tuple[Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) + smag_limit = _setup_smag_limit(diff_multfac_vn) + return diff_multfac_vn, smag_limit + + +@program +def setup_fields_for_initial_step( + k4: float, + hdiff_efdt_ratio: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], +): + _setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit)) + + +@field_operator +def _en_smag_fac_for_zero_nshift( + vect_a: Field[[KDim], float], + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, +) -> Field[[KDim], float]: + dz21 = hdiff_smag_z2 - hdiff_smag_z + alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 + df32 = hdiff_smag_fac3 - hdiff_smag_fac2 + df42 = hdiff_smag_fac4 - hdiff_smag_fac2 + dz32 = hdiff_smag_z3 - hdiff_smag_z2 + dz42 = hdiff_smag_z4 - hdiff_smag_z2 + + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (vect_a + vect_a(Koff[1])) + + dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z)) + dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2)) + enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) + return enh_smag_fac + + +@field_operator +def _init_diffusion_local_fields_for_regular_timestemp( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], +) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) + smag_limit = _setup_smag_limit(diff_multfac_vn) + enh_smag_fac = _en_smag_fac_for_zero_nshift( + vect_a, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + ) + return ( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ) + + +@program +def init_diffusion_local_fields_for_regular_timestep( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + enh_smag_fac: Field[[KDim], float], +): + _init_diffusion_local_fields_for_regular_timestemp( + k4, + dyn_substeps, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + vect_a, + out=( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ), + ) + + +def init_nabla2_factor_in_upper_damping_zone( + k_size: int, nrdmax: int32, nshift: int, physical_heights: np.ndarray +) -> Field[[KDim], float]: + """ + Calculate diff_multfac_n2w. + + numpy version, since gt4py does not allow non-constant indexing into fields + + Args + k_size: number of vertical levels + nrdmax: index of the level where rayleigh dampint starts + nshift: + physcial_heights: vector of physical heights [m] of the height levels + """ + # TODO(Magdalena): fix with as_offset in gt4py + + buffer = np.zeros(k_size) + buffer[1 : nrdmax + 1] = ( + 1.0 + / 12.0 + * ( + ( + physical_heights[1 + nshift : nrdmax + 1 + nshift] + - physical_heights[nshift + nrdmax + 1] + ) + / (physical_heights[1] - physical_heights[nshift + nrdmax + 1]) + ) + ** 4 + ) + return np_as_located_field(KDim)(buffer) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py similarity index 86% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py index 61560cb353..bceeecf447 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py @@ -14,12 +14,18 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import _calculate_nabla2_of_theta -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( _truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.model.atmosphere.dycore.update_theta_and_exner import _update_theta_and_exner +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + _update_theta_and_exner, +) from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py similarity index 85% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 333c21116d..36fded768e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -10,20 +10,20 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( _apply_nabla2_and_nabla4_global_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.model.atmosphere.dycore.calculate_nabla4 import _calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -91,7 +91,7 @@ def _apply_diffusion_to_vn( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], @@ -110,6 +110,10 @@ def apply_diffusion_to_vn( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, limited_area: bool, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_vn( u_vert, @@ -130,4 +134,8 @@ def apply_diffusion_to_vn( start_2nd_nudge_line_idx_e, limited_area, out=vn, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py similarity index 83% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 294eb2ed19..dd4f0a264f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,14 +14,16 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( _calculate_horizontal_gradients_for_turbulence, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import _calculate_nabla2_for_w +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( + _calculate_nabla2_for_w, +) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim @@ -91,6 +93,10 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( nrdmax: int32, interior_idx: int32, halo_idx: int32, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area, @@ -108,4 +114,8 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( interior_idx, halo_idx, out=(w, dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py similarity index 75% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py index f8d62fe062..9e181ee32b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType 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, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -39,5 +39,20 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w) + _apply_nabla2_to_w( + area, + z_nabla2_c, + geofac_n2s, + w, + diff_multfac_w, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py similarity index 74% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py index 42128fd8ad..657bdd4278 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim @@ -35,5 +35,19 @@ def apply_nabla2_to_w_in_upper_damping_layer( diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _apply_nabla2_to_w_in_upper_damping_layer( + w, + diff_multfac_n2w, + cell_area, + z_nabla2_c, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py similarity index 71% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py index e051aba455..6b497c45aa 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py @@ -12,23 +12,23 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, ) -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator def _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: @@ -43,12 +43,16 @@ def _calculate_diagnostic_quantities_for_turbulence( def calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec, @@ -58,4 +62,8 @@ def calculate_diagnostic_quantities_for_turbulence( diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py similarity index 72% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 49eec5e7cf..6a997cd21d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -12,12 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim @@ -43,7 +43,19 @@ def calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( - theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + theta_v, + theta_ref_mc, + thresh_tdiff, + kh_smag_e, + out=kh_smag_e, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py index 5a0a98d830..d174d72b08 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType 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, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -36,5 +36,18 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_horizontal_gradients_for_turbulence(w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy)) + _calculate_horizontal_gradients_for_turbulence( + w, + geofac_grg_x, + geofac_grg_y, + out=(dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py similarity index 67% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py index cca6712de1..d72cecd1d4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py @@ -12,10 +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 +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import _calculate_nabla2_of_theta +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim @@ -38,5 +42,19 @@ def calculate_nabla2_for_theta( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_nabla2_for_theta(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _calculate_nabla2_for_theta( + kh_smag_e, + inv_dual_edge_length, + theta_v, + geofac_div, + out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py similarity index 74% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py index cdba6c34f8..c40fab4415 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType 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, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -31,5 +31,17 @@ def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c) + _calculate_nabla2_for_w( + w, + geofac_n2s, + out=z_nabla2_c, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py similarity index 77% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index a90aaddea5..cb528e6fac 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -15,19 +15,19 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator def _temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s, axis=C2EDim) / diff_multfac_smag - div = neighbor_sum(vn(C2E) * geofac_div, axis=C2EDim) + kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div @@ -35,8 +35,8 @@ def _temporary_fields_for_turbulence_diagnostics( def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], kh_c: Field[[CellDim, KDim], float], div: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 92% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index fd189e6690..55eb44f2a8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -65,6 +65,10 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( vcoef: Field[[CECDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, @@ -76,4 +80,8 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( theta_v, z_temp, out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py index bc53f68979..b74d324851 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim @@ -39,5 +39,20 @@ def update_theta_and_exner( theta_v: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], rd_o_cvd: float, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _update_theta_and_exner(z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner)) + _update_theta_and_exner( + z_temp, + area, + theta_v, + exner, + rd_o_cvd, + out=(theta_v, exner), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py new file mode 100644 index 0000000000..7ba065d2c8 --- /dev/null +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -0,0 +1,24 @@ +# 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 + +""" +Initialize pytest. + +Workaround for pytest not discovering those configuration function, when they are added to the +diffusion_test/conftest.py folder +""" +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index b6b255adca..a29c25ea69 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py_common>=0.0.5" + "icon4py-common>=0.0.5" ] description = "ICON dynamical core." dynamic = ['version'] diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 7ee1151470..9e9274b362 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -10,24 +10,23 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pytest -from gt4py.next.program_processors.runners.roundtrip import executor -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh - -BACKENDS = {"embedded": executor} -MESHES = {"simple_mesh": SimpleMesh()} - - -@pytest.fixture( - ids=MESHES.keys(), - params=MESHES.values(), +from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + backend, + damping_height, + data_provider, + datapath, + download_ser_data, + grid_savepoint, + icon_grid, + linit, + mesh, + step_date_exit, + step_date_init, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, ) -def mesh(request): - return request.param - - -@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) -def backend(request): - return request.param diff --git a/model/common/.pre-commit-config.yaml b/model/common/.pre-commit-config.yaml index 3d45afc460..0a905ba43c 100644 --- a/model/common/.pre-commit-config.yaml +++ b/model/common/.pre-commit-config.yaml @@ -63,6 +63,16 @@ repos: rev: v1.3.0 hooks: - id: yesqa + additional_dependencies: + - flake8==4.0.1 + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments - repo: https://github.com/psf/black rev: '22.3.0' diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 0a90a8cb95..22a32d8ac9 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -22,7 +22,10 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1" + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "netcdf4>=1.6.0" ] description = "Shared code for the icon4py model." dynamic = ['version'] diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py new file mode 100644 index 0000000000..1c64d9f046 --- /dev/null +++ b/model/common/src/icon4py/model/common/constants.py @@ -0,0 +1,31 @@ +# 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 typing import Final + + +#: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), +#: see https://glossary.ametsoc.org/wiki/Gas_constant. +GAS_CONSTANT_DRY_AIR: Final[float] = 287.04 + +#: Specific heat at constant pressure [J/K/kg] +CPD: Final[float] = 1004.64 + +#: Gas constant for water vapor [J/K/kg], rv in ICON. +GAS_CONSTANT_WATER_VAPOR: Final[float] = 461.51 + +#: Av. gravitational acceleration [m/s^2] +GRAVITATIONAL_ACCELERATION: Final[float] = 9.8066 + +#: Default physics to dynamics time step ratio +DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0 diff --git a/model/common/src/icon4py/model/common/decomposition/__init__.py b/model/common/src/icon4py/model/common/decomposition/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/common/src/icon4py/model/common/decomposition/decomposed.py b/model/common/src/icon4py/model/common/decomposition/decomposed.py new file mode 100644 index 0000000000..23b278b65f --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/decomposed.py @@ -0,0 +1,251 @@ +# 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 +import logging +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +import ghex +import ghex.unstructured as unstructured +import numpy as np +import numpy.ma as ma +from gt4py.next import Dimension + +from icon4py.model.common.decomposition.parallel_setup import ProcessProperties +from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim +from icon4py.model.common.utils import builder + + +log = logging.getLogger(__name__) + + +class DomainDescriptorIdGenerator: + _counter = 0 + _roundtrips = 0 + + def __init__(self, parallel_props: ProcessProperties): + self._comm_size = parallel_props.comm_size + self._roundtrips = parallel_props.rank + self._base = self._roundtrips * self._comm_size + + def __call__(self): + next_id = self._base + self._counter + if self._counter + 1 >= self._comm_size: + self._roundtrips = self._roundtrips + self._comm_size + self._base = self._roundtrips * self._comm_size + self._counter = 0 + else: + self._counter = self._counter + 1 + return next_id + + +class DecompositionInfo: + class EntryType(int, Enum): + ALL = (0,) + OWNED = (1,) + HALO = 2 + + @builder + def with_dimension(self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray): + masked_global_index = ma.array(global_index, mask=owner_mask) + self._global_index[dim] = masked_global_index + + def __init__(self, klevels: int): + self._global_index = {} + self._klevels = klevels + + @property + def klevels(self): + return self._klevels + + def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return self._to_local_index(dim) + case DecompositionInfo.EntryType.HALO: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[~mask] + case DecompositionInfo.EntryType.OWNED: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[mask] + + def _to_local_index(self, dim): + data = ma.getdata(self._global_index[dim], subok=False) + assert data.ndim == 1 + return np.arange(data.shape[0]) + + def owner_mask(self, dim: Dimension) -> np.ndarray: + return self._global_index[dim].mask + + def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return ma.getdata(self._global_index[dim], subok=False) + case DecompositionInfo.EntryType.OWNED: + global_index = self._global_index[dim] + return ma.getdata(global_index[global_index.mask]) + case DecompositionInfo.EntryType.HALO: + global_index = self._global_index[dim] + return ma.getdata(global_index[~global_index.mask]) + case _: + raise NotImplementedError() + + +class ExchangeResult(Protocol): + def wait(self): + ... + + def is_ready(self) -> bool: + ... + + +class ExchangeRuntime(Protocol): + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + ... + + def get_size(self): + ... + + def my_rank(self): + ... + + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +def create_exchange(props: ProcessProperties, decomp_info: DecompositionInfo) -> ExchangeRuntime: + """ + Create an Exchange depending on the runtime size. + + Depending on the number of processor a SingleNode version is returned or a GHEX context created and a Multinode returned. + """ + if props.comm_size > 1: + return GHexMultiNode(props, decomp_info) + else: + return SingleNode() + + +@dataclass +class SingleNode: + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + return SingleNodeResult() + + def my_rank(self): + return 0 + + def get_size(self): + return 1 + + +class SingleNodeResult: + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +class GHexMultiNode: + def __init__(self, props: ProcessProperties, domain_decomposition: DecompositionInfo): + self._context = ghex.context(ghex.mpi_comm(props.comm), True) + self._domain_id_gen = DomainDescriptorIdGenerator(props) + self._decomposition_info = domain_decomposition + self._domain_descriptors = { + CellDim: self._create_domain_descriptor( + CellDim, + ), + VertexDim: self._create_domain_descriptor( + VertexDim, + ), + EdgeDim: self._create_domain_descriptor(EdgeDim), + } + log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") + + self._patterns = { + CellDim: self._create_pattern(CellDim), + VertexDim: self._create_pattern(VertexDim), + EdgeDim: self._create_pattern(EdgeDim), + } + log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") + self._comm = unstructured.make_co(self._context) + log.info("communication object initialized") + + def _domain_descriptor_info(self, descr): + return f" domain_descriptor=[id='{descr.domain_id()}', size='{descr.size()}', inner_size='{descr.inner_size()}' (halo size='{descr.size() - descr.inner_size()}')" + + def get_size(self): + return self._context.size() + + def my_rank(self): + return self._context.rank() + + def _create_domain_descriptor(self, dim: Dimension): + all_global = self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) + local_halo = self._decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) + # first arg is the domain ID which builds up an MPI Tag. + # if those ids are not different for all domain descriptors the system might deadlock + # if two parallel exchanges with the same domain id are done + domain_desc = unstructured.domain_descriptor( + self._domain_id_gen(), all_global.tolist(), local_halo.tolist() + ) + log.debug( + f"domain descriptor for dim='{dim.value}' with properties {self._domain_descriptor_info(domain_desc)} created" + ) + return domain_desc + + def _create_pattern(self, horizontal_dim: Dimension): + assert horizontal_dim.kind == DimensionKind.HORIZONTAL + + global_halo_idx = self._decomposition_info.global_index( + horizontal_dim, DecompositionInfo.EntryType.HALO + ) + halo_generator = unstructured.halo_generator_with_gids(global_halo_idx) + log.debug(f"halo generator for dim='{horizontal_dim.value}' created") + pattern = unstructured.make_pattern( + self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] + ) + log.debug( + f"pattern for dim='{horizontal_dim.value}' and {self._domain_descriptor_info(self._domain_descriptors[horizontal_dim])} created" + ) + return pattern + + def exchange(self, dim: Dimension, *fields: tuple): + assert dim in [CellDim, EdgeDim, VertexDim] + pattern = self._patterns[dim] + assert pattern is not None, f"pattern for {dim.value} not found" + domain_descriptor = self._domain_descriptors[dim] + assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" + applied_patterns = [ + pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields + ] + handle = self._comm.exchange(applied_patterns) + log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") + return MultiNodeResult(handle, applied_patterns) + + +@dataclass +class MultiNodeResult: + handle: ... + pattern_refs: ... + + def wait(self): + self.handle.wait() + del self.pattern_refs + + def is_ready(self) -> bool: + return self.handle.is_ready() diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py new file mode 100644 index 0000000000..11d5eaeb6b --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -0,0 +1,98 @@ +# 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 +import functools +import logging +from dataclasses import dataclass +from typing import Optional, Union + +import mpi4py +from mpi4py.MPI import Comm + + +mpi4py.rc.initialize = False + +CommId = Union[int, Comm, None] +log = logging.getLogger(__name__) + + +def get_processor_properties(with_mpi=False, comm_id: CommId = None): + def _get_current_comm_or_comm_world(comm_id: CommId) -> Comm: + if isinstance(comm_id, int): + comm = Comm.f2py(comm_id) + elif isinstance(comm_id, Comm): + comm = comm_id + else: + comm = mpi4py.MPI.COMM_WORLD + return comm + + if with_mpi: + init_mpi() + current_comm = _get_current_comm_or_comm_world(comm_id) + return ProcessProperties.from_mpi_comm(current_comm) + else: + return ProcessProperties.from_single_node() + + +def init_mpi(): + from mpi4py import MPI + + if not MPI.Is_initialized(): + log.info("initializing MPI") + MPI.Init() + + +def finalize_mpi(): + from mpi4py import MPI + + if not MPI.Is_finalized(): + log.info("finalizing MPI") + MPI.Finalize() + + +@dataclass(frozen=True) +class ProcessProperties: + comm: Optional[mpi4py.MPI.Comm] = None + + @functools.cached_property + def rank(self): + return self.comm.Get_rank() if self.comm else 0 + + @functools.cached_property + def comm_name(self): + return self.comm.Get_name() if self.comm else "" + + @functools.cached_property + def comm_size(self): + return self.comm.Get_size() if self.comm else 1 + + @classmethod + def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): + return ProcessProperties(comm) + + @classmethod + def from_single_node(cls): + return ProcessProperties() + + +class ParallelLogger(logging.Filter): + def __init__(self, process_properties: ProcessProperties = None): + super().__init__() + self._rank_info = "" + if process_properties and process_properties.comm_size > 1: + self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " + + def filter( # noqa: A003 # overwriting logging.Filter.filter() + self, record: logging.LogRecord + ) -> bool: + record.rank = self._rank_info + return True diff --git a/model/common/src/icon4py/model/common/dimension.py b/model/common/src/icon4py/model/common/dimension.py index 8265dcfa6d..49c74a68bb 100644 --- a/model/common/src/icon4py/model/common/dimension.py +++ b/model/common/src/icon4py/model/common/dimension.py @@ -30,7 +30,9 @@ E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) V2CDim = Dimension("V2C", DimensionKind.LOCAL) +C2VDim = Dimension("C2V", DimensionKind.LOCAL) V2EDim = Dimension("V2E", DimensionKind.LOCAL) +V2E2VDim = Dimension("V2E2V", DimensionKind.LOCAL) E2C2VDim = Dimension("E2C2V", DimensionKind.LOCAL) C2E2CODim = Dimension("C2E2CO", DimensionKind.LOCAL) E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL) @@ -40,6 +42,7 @@ E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim)) C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim)) V2C = FieldOffset("V2C", source=CellDim, target=(VertexDim, V2CDim)) +C2V = FieldOffset("C2V", source=VertexDim, target=(CellDim, C2VDim)) V2E = FieldOffset("V2E", source=EdgeDim, target=(VertexDim, V2EDim)) E2V = FieldOffset("E2V", source=VertexDim, target=(EdgeDim, E2VDim)) C2CE = FieldOffset("C2CE", source=CEDim, target=(CellDim, C2EDim)) diff --git a/model/common/src/icon4py/model/common/grid/__init__.py b/model/common/src/icon4py/model/common/grid/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py new file mode 100644 index 0000000000..79382c3542 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -0,0 +1,405 @@ +# 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 +import dataclasses +import logging +from enum import Enum +from typing import Optional +from uuid import UUID + +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from netCDF4 import Dataset + +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + C2VDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + EdgeDim, + V2CDim, + V2E2VDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize + + +class GridFileName(str, Enum): + pass + + +@dataclasses.dataclass +class GridFileField: + name: GridFileName + shape: tuple[int, ...] + + +def _validate_shape(data: np.array, field_definition: GridFileField): + if data.shape != field_definition.shape: + raise IconGridError( + f"invalid grid file field {field_definition.name} does not have dimension {field_definition.shape}" + ) + + +class GridFile: + """Represent and ICON netcdf grid file.""" + + INVALID_INDEX = -1 + + class PropertyName(GridFileName): + GRID_ID = "uuidOfHGrid" + PARENT_GRID_ID = "uuidOfParHGrid" + + class OffsetName(GridFileName): + """Names for connectivities used in the grid file.""" + + # e2c2e/e2c2eO: diamond edges (including origin) not present in grid file-> calculate? + # e2c2v: diamond vertices: not present in grid file -> constructed from e2c and c2v + + #: name of C2E2C connectivity in grid file: dims(nv=3, cell) + C2E2C = "neighbor_cell_index" + + #: name of V2E2V connectivity in gridfile: dims(ne=6, vertex), + #: all vertices of a pentagon/hexagon, same as V2C2V + V2E2V = "vertices_of_vertex" # does not exist in simple_mesh.py + + #: name of V2E dimension in grid file: dims(ne=6, vertex) + V2E = "edges_of_vertex" + + #: name fo V2C connectivity in grid file: dims(ne=6, vertex) + V2C = "cells_of_vertex" + + #: name of E2V connectivity in grid file: dims(nc=2, edge) + E2V = "edge_vertices" + + #: name of C2V connectivity in grid file: dims(nv=3, cell) + C2V = "vertex_of_cell" # does not exist in simple_mesh.py + + #: name of E2C connectivity in grid file: dims(nc=2, edge) + E2C = "adjacent_cell_of_edge" + + #: name of C2E connectivity in grid file: dims(nv=3, cell) + C2E = "edge_of_cell" + + class DimensionName(GridFileName): + """Dimension values (sizes) used in grid file.""" + + #: number of vertices + VERTEX_NAME = "vertex" + + #: number of edges + EDGE_NAME = "edge" + + #: number of cells + CELL_NAME = "cell" + + #: number of edges in a diamond: 4 + DIAMOND_EDGE_SIZE = "no" + + #: number of edges/cells neibhboring one vertex: 6 (for regular, non pentagons) + NEIGHBORS_TO_VERTEX_SIZE = "ne" + + #: number of cells edges, vertices and cells neighboring a cell: 3 + NEIGHBORS_TO_CELL_SIZE = "nv" + + #: number of vertices/cells neighboring an edge: 2 + NEIGHBORS_TO_EDGE_SIZE = "nc" + + #: number of child domains (for nesting) + MAX_CHILD_DOMAINS = "max_chdom" + + #: Grid refinement: maximal number in grid-refinement (refin_ctl) array for each dimension + CELL_GRF = "cell_grf" + EDGE_GRF = "edge_grf" + VERTEX_GRF = "vert_grf" + + class GridRefinementName(GridFileName): + """Names of arrays in grid file defining the grid control, definition of boundaries layers, start and end indices of horizontal zones.""" + + #: refine control value of cell indices + CONTROL_CELLS = "refin_c_ctrl" + + #: refine control value of edge indices + CONTROL_EDGES = "refin_e_ctrl" + + #: refine control value of vertex indices + CONTROL_VERTICES = "refin_v_ctrl" + + #: start indices of horizontal grid zones for cell fields + START_INDEX_CELLS = "start_idx_c" + + #: start indices of horizontal grid zones for edge fields + START_INDEX_EDGES = "start_idx_e" + + #: start indices of horizontal grid zones for vertex fields + START_INDEX_VERTICES = "start_idx_v" + + #: end indices of horizontal grid zones for cell fields + END_INDEX_CELLS = "end_idx_c" + + #: end indices of horizontal grid zones for edge fields + END_INDEX_EDGES = "end_idx_e" + + #: end indices of horizontal grid zones for vertex fields + END_INDEX_VERTICES = "end_idx_v" + + def __init__(self, dataset: Dataset): + self._dataset = dataset + self._log = logging.getLogger(__name__) + + def dimension(self, name: GridFileName) -> int: + return self._dataset.dimensions[name].size + + def int_field(self, name: GridFileName, transpose=True, dtype=np.int32) -> np.ndarray: + try: + nc_variable = self._dataset.variables[name] + + self._log.debug(f"reading {name}: {nc_variable}") + data = nc_variable[:] + data = np.array(data, dtype=dtype) + return np.transpose(data) if transpose else data + except KeyError: + msg = f"{name} does not exist in dataset" + self._log.warning(msg) + raise IconGridError(msg) + + +class IconGridError(RuntimeError): + pass + + +class IndexTransformation: + def get_offset_for_index_field( + self, + array: np.ndarray, + ): + return np.zeros(array.shape, dtype=np.int32) + + +class ToGt4PyTransformation(IndexTransformation): + def get_offset_for_index_field(self, array: np.ndarray): + """ + Calculate the index offset needed for usage with python. + + Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for + INVALID values which are marked with -1 in the grid file and are kept such. + """ + return np.where(array == GridFile.INVALID_INDEX, 0, -1) + + +class GridManager: + """ + Read ICON grid file and set up IconGrid. + + Reads an ICON grid file and extracts connectivity arrays and start-, end-indices for horizontal + domain boundaries. Provides an IconGrid instance for further usage. + """ + + def __init__( + self, + transformation: IndexTransformation, + grid_file: str, + config: VerticalGridSize, + ): + self._log = logging.getLogger(__name__) + self._transformation = transformation + self._config = config + self._grid: Optional[IconGrid] = None + self._file_name = grid_file + + def __call__(self): + dataset = self._read_gridfile(self._file_name) + _, grid = self._constuct_grid(dataset) + self._grid = grid + + def _read_gridfile(self, fname: str) -> Dataset: + try: + dataset = Dataset(self._file_name, "r", format="NETCDF4") + self._log.debug(dataset) + return dataset + except FileNotFoundError: + self._log.error(f"gridfile {fname} not found, aborting") + exit(1) + + def _read_grid_refinement_information(self, dataset): + _CHILD_DOM = 0 + reader = GridFile(dataset) + + refin_ctrl = { + CellDim: reader.int_field(GridFile.GridRefinementName.CONTROL_CELLS), + EdgeDim: reader.int_field(GridFile.GridRefinementName.CONTROL_EDGES), + VertexDim: reader.int_field(GridFile.GridRefinementName.CONTROL_VERTICES), + } + refin_ctrl_max = { + CellDim: reader.dimension(GridFile.DimensionName.CELL_GRF), + EdgeDim: reader.dimension(GridFile.DimensionName.EDGE_GRF), + VertexDim: reader.dimension(GridFile.DimensionName.VERTEX_GRF), + } + start_indices = { + CellDim: self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_EDGES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_VERTICES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + } + end_indices = { + CellDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_CELLS, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_EDGES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_VERTICES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + } + + return start_indices, end_indices, refin_ctrl, refin_ctrl_max + + def get_grid(self): + return self._grid + + def _get_index(self, dim: Dimension, start_marker: int, index_dict): + if dim.kind != DimensionKind.HORIZONTAL: + msg = f"getting start index in horizontal domain with non - horizontal dimension {dim}" + self._log.warning(msg) + raise IconGridError(msg) + try: + return index_dict[dim][start_marker] + except KeyError: + msg = f"start, end indices for dimension {dim} not present" + self._log.error(msg) + raise IconGridError(msg) + + def _constuct_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: + grid_id = UUID(dataset.getncattr(GridFile.PropertyName.GRID_ID)) + return grid_id, self._from_grid_dataset(dataset) + + def get_size(self, dim: Dimension): + if dim == VertexDim: + return self._grid.config.num_vertices + elif dim == CellDim: + return self._grid.config.num_cells + elif dim == EdgeDim: + return self._grid.config.num_edges + else: + self._log.warning(f"cannot determine size of unknown dimension {dim}") + raise IconGridError(f"Unknown dimension {dim}") + + def _get_index_field( + self, + reader, + field: GridFileName, + transpose=True, + apply_offset=True, + dtype=np.int32, + ): + field = reader.int_field(field, transpose=transpose, dtype=dtype) + if apply_offset: + field = field + self._transformation.get_offset_for_index_field(field) + return field + + def _from_grid_dataset(self, dataset: Dataset) -> IconGrid: + reader = GridFile(dataset) + num_cells = reader.dimension(GridFile.DimensionName.CELL_NAME) + num_edges = reader.dimension(GridFile.DimensionName.EDGE_NAME) + num_vertices = reader.dimension(GridFile.DimensionName.VERTEX_NAME) + + grid_size = HorizontalGridSize( + num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells + ) + c2e = self._get_index_field(reader, GridFile.OffsetName.C2E) + + e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) + c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) + e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) + + e2c2v = self._construct_diamond_array(c2v, e2c) + + v2c = self._get_index_field(reader, GridFile.OffsetName.V2C) + v2e = self._get_index_field(reader, GridFile.OffsetName.V2E) + v2e2v = self._get_index_field(reader, GridFile.OffsetName.V2E2V) + c2e2c = self._get_index_field(reader, GridFile.OffsetName.C2E2C) + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + ( + start_indices, + end_indices, + refine_ctrl, + refine_ctrl_max, + ) = self._read_grid_refinement_information(dataset) + + config = GridConfig( + horizontal_config=grid_size, + vertical_config=self._config, + ) + icon_grid = ( + IconGrid() + .with_config(config) + .with_connectivities( + { + C2EDim: c2e, + E2CDim: e2c, + E2VDim: e2v, + V2EDim: v2e, + V2CDim: v2c, + C2VDim: c2v, + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + E2C2VDim: e2c2v, + V2E2VDim: v2e2v, + } + ) + .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim]) + .with_start_end_indices(EdgeDim, start_indices[EdgeDim], end_indices[EdgeDim]) + .with_start_end_indices(VertexDim, start_indices[VertexDim], end_indices[VertexDim]) + ) + + return icon_grid + + def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): + dummy_c2v = np.append( + c2v, + GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), + axis=0, + ) + expanded = dummy_c2v[e2c[:, :], :] + sh = expanded.shape + flattened = expanded.reshape(sh[0], sh[1] * sh[2]) + return np.apply_along_axis(np.unique, 1, flattened) diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py new file mode 100644 index 0000000000..a68c95d1bd --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -0,0 +1,269 @@ +# 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 dataclasses import dataclass +from typing import Final + +from gt4py.next.common import Dimension, Field + +from icon4py.model.common.dimension import CellDim, ECVDim, EdgeDim, VertexDim + + +class HorizontalMarkerIndex: + """ + Handles constants indexing into the start_index and end_index fields. + + ICON uses a double indexing scheme for field indices marking the start and end of special + grid zone: The constants defined here (from mo_impl_constants.f90 and mo_impl_constants_grf.f90) + are the indices that are used to index into the start_idx and end_idx arrays + provided by the grid file where for each dimension the start index of the horizontal + "zones" are defined: + f.ex. an inlined access of the field F: Field[[CellDim], double] at the starting point of the lateral boundary zone would be + + F[start_idx_c[_LATERAL_BOUNDARY_CELLS] + + ICON uses a custom index range from [ICON_INDEX_OFFSET... ] such that the index 0 marks the + internal entities for _all_ dimensions (Cell, Edge, Vertex) that is why we define these + additional INDEX_OFFSETs here in order to swap back to a 0 base python array. + """ + + NUM_GHOST_ROWS: Final[int] = 2 + # values from mo_impl_constants.f90 + _ICON_INDEX_OFFSET_CELLS: Final[int] = 8 + _GRF_BOUNDARY_WIDTH_CELL: Final[int] = 4 + _MIN_RL_CELL_INT: Final[int] = -4 + _MIN_RL_CELL: Final[int] = _MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS + _MAX_RL_CELL: Final[int] = 5 + + _ICON_INDEX_OFFSET_VERTEX: Final[int] = 7 + _MIN_RL_VERTEX_INT: Final[int] = _MIN_RL_CELL_INT + _MIN_RL_VERTEX: Final[int] = _MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1) + _MAX_RL_VERTEX: Final[int] = _MAX_RL_CELL + + _ICON_INDEX_OFFSET_EDGES: Final[int] = 13 + _GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 + _MIN_RL_EDGE_INT: Final[int] = 2 * _MIN_RL_CELL_INT + _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) + _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL + + _LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES + _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES + _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES + _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES + _END_EDGES: Final[int] = 0 + + _LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS + _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS + _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS + _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS + _END_CELLS: Final[int] = 0 + + _LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX + _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX + _NUDGING_VERTICES: Final[int] = 0 + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX + _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX + _END_VERTICES: Final[int] = 0 + + _lateral_boundary = { + CellDim: _LATERAL_BOUNDARY_CELLS, + EdgeDim: _LATERAL_BOUNDARY_EDGES, + VertexDim: _LATERAL_BOUNDARY_VERTICES, + } + _local = { + CellDim: _LOCAL_CELLS, + EdgeDim: _LOCAL_EDGES, + VertexDim: _LOCAL_VERTICES, + } + _halo = { + CellDim: _HALO_CELLS, + EdgeDim: _HALO_EDGES, + VertexDim: _HALO_VERTICES, + } + _interior = { + CellDim: _INTERIOR_CELLS, + EdgeDim: _INTERIOR_EDGES, + VertexDim: _INTERIOR_VERTICES, + } + _nudging = { + CellDim: _NUDGING_CELLS, + EdgeDim: _NUDGING_EDGES, + VertexDim: _NUDGING_VERTICES, + } + _end = { + CellDim: _END_CELLS, + EdgeDim: _END_EDGES, + VertexDim: _END_VERTICES, + } + + @classmethod + def lateral_boundary(cls, dim: Dimension) -> int: + """Indicate lateral boundary. + + These points correspond to the sorted points in ICON, the marker can be incremented in order + to access higher order boundary lines + """ + return cls._lateral_boundary[dim] + + @classmethod + def local(cls, dim: Dimension) -> int: + """Indicate points that are owned by the processing unit, i.e. no halo points.""" + return cls._local[dim] + + @classmethod + def halo(cls, dim: Dimension) -> int: + return cls._halo[dim] + + @classmethod + def nudging(cls, dim: Dimension) -> int: + """Indicate the nudging zone.""" + return cls._nudging[dim] + + @classmethod + def interior(cls, dim: Dimension) -> int: + """Indicate interior i.e. unordered prognostic cells in ICON.""" + return cls._interior[dim] + + @classmethod + def end(cls, dim: Dimension) -> int: + return cls._end[dim] + + +@dataclass(frozen=True) +class HorizontalGridSize: + num_vertices: int + num_edges: int + num_cells: int + + +# TODO(Magdalena): allow initialization with only partial values +# (becomes tedious for testing otherwise): hence this should +# that should not be a data class +class EdgeParams: + def __init__( + self, + tangent_orientation=None, + primal_edge_lengths=None, + inverse_primal_edge_lengths=None, + dual_edge_lengths=None, + inverse_dual_edge_lengths=None, + inverse_vertex_vertex_lengths=None, + primal_normal_vert_x=None, + primal_normal_vert_y=None, + dual_normal_vert_x=None, + dual_normal_vert_y=None, + edge_areas=None, + ): + + self.tangent_orientation: Field[[EdgeDim], float] = tangent_orientation + r""" + Orientation of vector product of the edge and the adjacent cell centers + v3 + / \ + / \ + / c1 \ + / | \ + v1---|--->v2 + \ | / + \ v / + \ c2 / + \ / + v4 + +1 or -1 depending on whether the vector product of + (v2-v1) x (c2-c1) points outside (+) or inside (-) the sphere + + defined in ICON in mo_model_domain.f90:t_grid_edges%tangent_orientation + """ + + self.primal_edge_lengths: Field[[EdgeDim], float] = primal_edge_lengths + """ + Length of the triangle edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length + """ + + self.inverse_primal_edge_lengths: Field[[EdgeDim], float] = inverse_primal_edge_lengths + """ + Inverse of the triangle edge length: 1.0/primal_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_primal_edge_length + """ + + self.dual_edge_lengths: Field[[EdgeDim], float] = dual_edge_lengths + """ + Length of the hexagon/pentagon edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length + """ + + self.inverse_dual_edge_lengths: Field[[EdgeDim], float] = inverse_dual_edge_lengths + """ + Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length + """ + + self.inverse_vertex_vertex_lengths: Field[[EdgeDim], float] = inverse_vertex_vertex_lengths + r""" + Inverse distance between outer vertices of adjacent cells. + + v1-------- + | /| + | / | + | e | + | / | + |/ | + --------v2 + + inverse_vertex_vertex_length(e) = 1.0/|v2-v1| + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length + """ + + self.primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + primal_normal_vert_x, + primal_normal_vert_y, + ) + """ + Normal of the triangle edge, projected onto the location of the vertices + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert + """ + + self.dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + dual_normal_vert_x, + dual_normal_vert_y, + ) + """ + Tangent to the triangle edge, projected onto the location of vertices. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_normal_vert + """ + + self.edge_areas: Field[[EdgeDim], float] = edge_areas + """ + Area of the quadrilateral (two triangles) adjacent to the edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%area_edge + """ + + +@dataclass(frozen=True) +class CellParams: + area: Field[[CellDim], float] + """ + Area of a cell. + + defined int ICON in mo_model_domain.f90:t_grid_cells%area + """ diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py new file mode 100644 index 0000000000..f514ca1285 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -0,0 +1,197 @@ +# 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 dataclasses import dataclass +from typing import Dict + +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import NeighborTableOffsetProvider +from typing_extensions import deprecated + +from icon4py.model.common.dimension import CECDim, CEDim, CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.horizontal import HorizontalGridSize +from icon4py.model.common.grid.vertical import VerticalGridSize +from icon4py.model.common.utils import builder + + +@dataclass(frozen=True) +class GridConfig: + horizontal_config: HorizontalGridSize + vertical_config: VerticalGridSize + limited_area: bool = True + n_shift_total: int = 0 + + @property + def num_k_levels(self): + return self.vertical_config.num_lev + + @property + def num_vertices(self): + return self.horizontal_config.num_vertices + + @property + def num_edges(self): + return self.horizontal_config.num_edges + + @property + def num_cells(self): + return self.horizontal_config.num_cells + + +class IconGrid: + def __init__(self): + self.config: GridConfig = None + + self.start_indices = {} + self.end_indices = {} + self.connectivities: Dict[str, np.ndarray] = {} + self.size: Dict[Dimension, int] = {} + + def _update_size(self, config: GridConfig): + self.size[VertexDim] = config.num_vertices + self.size[CellDim] = config.num_cells + self.size[EdgeDim] = config.num_edges + self.size[KDim] = config.num_k_levels + + @builder + def with_config(self, config: GridConfig): + self.config = config + self._update_size(config) + + @builder + def with_start_end_indices( + self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray + ): + self.start_indices[dim] = start_indices.astype(int32) + self.end_indices[dim] = end_indices.astype(int32) + + @builder + def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): + self.connectivities.update( + {d.value.lower(): k.astype(int) for d, k in connectivity.items()} + ) + self.size.update({d: t.shape[1] for d, t in connectivity.items()}) + + def limited_area(self): + # defined in mo_grid_nml.f90 + return self.config.limited_area + + def n_lev(self): + return self.config.num_k_levels if self.config else 0 + + def num_cells(self): + return self.config.num_cells if self.config else 0 + + def num_vertices(self): + return self.config.num_vertices if self.config else 0 + + def num_edges(self): + return self.config.num_edges + + @deprecated( + "use get_start_index and get_end_index instead, - should be removed after merge of solve_nonhydro" + ) + def get_indices_from_to( + self, dim: Dimension, start_marker: int, end_marker: int + ) -> tuple[int32, int32]: + """ + Use to specify domains of a field for field_operator. + + For a given dimension, returns the start and end index if a + horizontal region in a field given by the markers. + + field operators will then run from start of the region given by the + start_marker to the end of the region given by the end_marker + """ + if dim.kind != DimensionKind.HORIZONTAL: + raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) + return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] + + def get_start_index(self, dim: Dimension, marker: int) -> int32: + """ + Use to specify lower end of domains of a field for field_operators. + + For a given dimension, returns the start index of the + horizontal region in a field given by the marker. + """ + return self.start_indices[dim][marker] + + def get_end_index(self, dim: Dimension, marker: int) -> int32: + """ + Use to specify upper end of domains of a field for field_operators. + + For a given dimension, returns the end index of the + horizontal region in a field given by the marker. + """ + return self.end_indices[dim][marker] + + def get_c2e_connectivity(self): + table = self.connectivities["c2e"] + return NeighborTableOffsetProvider(table, CellDim, EdgeDim, table.shape[1]) + + def get_e2c_connectivity(self): + table = self.connectivities["e2c"] + return NeighborTableOffsetProvider(table, EdgeDim, CellDim, table.shape[1]) + + def get_e2v_connectivity(self): + table = self.connectivities["e2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) + + def get_c2e2c_connectivity(self): + table = self.connectivities["c2e2c"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + def get_c2e2co_connectivity(self): + table = self.connectivities["c2e2co"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + def get_e2c2v_connectivity(self): + table = self.connectivities["e2c2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) + + def get_v2e_connectivity(self): + table = self.connectivities["v2e"] + return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + + def get_v2c_connectivity(self): + table = self.connectivities["v2c"] + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) + + def get_c2v_connectivity(self): + table = self.connectivities["c2v"] + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) + + def get_e2ecv_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["e2c2v"].shape, EdgeDim, ECVDim + ) + + def _neighbortable_offset_provider_for_1d_sparse_fields( + self, + old_shape: tuple[int, int], + origin_axis: Dimension, + neighbor_axis: Dimension, + ): + table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(table, origin_axis, neighbor_axis, table.shape[1]) + + def get_c2cec_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e2c"].shape, CellDim, CECDim + ) + + def get_c2ce_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e"].shape, CellDim, CEDim + ) diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py new file mode 100644 index 0000000000..061e6a786a --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -0,0 +1,56 @@ +# 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 dataclasses import Field, dataclass, field +from typing import Final + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.common.dimension import KDim + + +@dataclass(frozen=True) +class VerticalGridSize: + num_lev: int + + +@dataclass(frozen=True) +class VerticalModelParams: + """ + Contains vertical physical parameters defined on the grid. + + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml + """ + + vct_a: Field[[KDim], float] + rayleigh_damping_height: Final[float] + index_of_damping_layer: Final[int32] = field(init=False) + + def __post_init__(self): + object.__setattr__( + self, + "index_of_damping_layer", + self._determine_damping_height_index( + np.asarray(self.vct_a), self.rayleigh_damping_height + ), + ) + + @classmethod + def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): + return int32(np.argmax(np.where(vct_a >= damping_height))) + + @property + def physical_heights(self) -> Field[[KDim], float]: + return self.vct_a diff --git a/model/common/src/icon4py/model/common/interpolation/__init__.py b/model/common/src/icon4py/model/common/interpolation/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py new file mode 100644 index 0000000000..cd7fa4ab78 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -0,0 +1,51 @@ +# 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 +# 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 + +import numpy as np + + +def compute_c_lin_e( + edge_cell_length: np.array, + inv_dual_edge_length: np.array, + owner_mask: np.array, + second_boundary_layer_start_index: np.int32, +) -> np.array: + """ + Compute E2C average inverse distance. + + Args: + edge_cell_length: numpy array, representing a Field[[EdgeDim, E2CDim], float] + inv_dual_edge_length: inverse dual edge length, numpy array representing a Field[[EdgeDim], float] + owner_mask: numpy array, representing a Field[[EdgeDim], bool]boolean field, True for all edges owned by this compute node + second_boundary_layer_start_index: start index of the 2nd boundary line: c_lin_e is not calculated for the first boundary layer + + Returns: c_lin_e: numpy array representing Field[[EdgeDim, E2CDim], float] + + """ + c_lin_e_ = edge_cell_length[:, 1] * inv_dual_edge_length + c_lin_e = np.transpose(np.vstack((c_lin_e_, (1.0 - c_lin_e_)))) + c_lin_e[0:second_boundary_layer_start_index, :] = 0.0 + mask = np.transpose(np.tile(owner_mask, (2, 1))) + return np.where(mask, c_lin_e, 0.0) diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py index 0a9bf2d689..bf5bb4c3c8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType 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, int32, neighbor_sum from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -36,5 +36,18 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) + _mo_intp_rbf_rbf_vec_interpol_vertex( + p_e_in, + ptr_coeff_1, + ptr_coeff_2, + out=(p_u_out, p_v_out), + domain={ + VertexDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/test_utils/data_handling.py b/model/common/src/icon4py/model/common/test_utils/data_handling.py new file mode 100644 index 0000000000..7a8a49c345 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/data_handling.py @@ -0,0 +1,30 @@ +# 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 + +import tarfile +from pathlib import Path + +import wget + + +def download_and_extract(uri: str, local_path: Path, data_file: str) -> None: + local_path.mkdir(parents=True, exist_ok=True) + if not any(local_path.iterdir()): + print(f"directory {local_path} is empty: downloading data from {uri} and extracting") + wget.download(uri, out=data_file) + # extract downloaded file + if not tarfile.is_tarfile(data_file): + raise NotImplementedError(f"{data_file} needs to be a valid tar file") + with tarfile.open(data_file, mode="r:*") as tf: + tf.extractall(path=local_path) + Path(data_file).unlink(missing_ok=True) diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py new file mode 100644 index 0000000000..c6a52f0c94 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -0,0 +1,183 @@ +# 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 pathlib import Path + +import pytest +from gt4py.next.program_processors.runners.roundtrip import executor + +from ..decomposition.parallel_setup import get_processor_properties +from .data_handling import download_and_extract +from .serialbox_utils import IconSerialDataProvider +from .simple_mesh import SimpleMesh + + +test_utils = Path(__file__).parent +model = test_utils.parent.parent +common = model.parent.parent.parent.parent +base_path = common.parent.joinpath("testdata") + +data_uris = { + 1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download", + 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", + 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", +} + +ser_data_basepath = base_path.joinpath("ser_icondata") + + +@pytest.fixture(params=[False], scope="session") +def processor_props(request): + with_mpi = request.param + return get_processor_properties(with_mpi=with_mpi) + + +@pytest.fixture(scope="session") +def ranked_data_path(processor_props): + return ser_data_basepath.absolute().joinpath(f"mpitask{processor_props.comm_size}") + + +@pytest.fixture(scope="session") +def datapath(ranked_data_path): + return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") + + +@pytest.fixture(scope="session") +def download_ser_data(request, processor_props, ranked_data_path, pytestconfig): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + if not pytestconfig.getoption("datatest"): + pytest.skip("not running datatest marked tests") + + try: + uri = data_uris[processor_props.comm_size] + + data_file = ranked_data_path.joinpath( + f"mch_ch_r04b09_dsl_mpitask{processor_props.comm_size}.tar.gz" + ).name + if processor_props.rank == 0: + download_and_extract(uri, ranked_data_path, data_file) + if processor_props.comm: + processor_props.comm.barrier() + except KeyError: + raise AssertionError( + f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" + ) + + +@pytest.fixture(scope="session") +def data_provider(download_ser_data, datapath, processor_props) -> IconSerialDataProvider: + return IconSerialDataProvider( + fname_prefix="icon_pydycore", + path=str(datapath), + mpi_rank=processor_props.rank, + do_print=True, + ) + + +@pytest.fixture +def grid_savepoint(data_provider): + return data_provider.from_savepoint_grid() + + +@pytest.fixture +def icon_grid(grid_savepoint): + """ + Load the icon grid from an ICON savepoint. + + Uses the special grid_savepoint that contains data from p_patch + """ + return grid_savepoint.construct_icon_grid() + + +@pytest.fixture +def decomposition_info(data_provider): + return data_provider.from_savepoint_grid().construct_decomposition_info() + + +@pytest.fixture +def damping_height(): + return 12500 + + +@pytest.fixture +def ndyn_substeps(): + """ + Return number of dynamical substeps. + + Serialized data uses a reduced number (2 instead of the default 5) in order to reduce the amount + of data generated. + """ + return 2 + + +@pytest.fixture +def linit(): + """ + Set the 'linit' flag for the ICON diffusion data savepoint. + + Defaults to False + """ + return False + + +@pytest.fixture +def step_date_init(): + """ + Set the step date for the loaded ICON time stamp at start of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def step_date_exit(): + """ + Set the step date for the loaded ICON time stamp at the end of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def interpolation_savepoint(data_provider): # F811 + """Load data from ICON interplation state savepoint.""" + return data_provider.from_interpolation_savepoint() + + +@pytest.fixture +def metrics_savepoint(data_provider): # F811 + """Load data from ICON mestric state savepoint.""" + return data_provider.from_metrics_savepoint() + + +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py new file mode 100644 index 0000000000..6dbe56db6d --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -0,0 +1,22 @@ +# 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 + + +import pytest + +from icon4py.model.common.decomposition.parallel_setup import ProcessProperties + + +def check_comm_size(props: ProcessProperties, sizes=(1, 2, 4)): + if props.comm_size not in sizes: + pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py new file mode 100644 index 0000000000..de4425c8aa --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/pytest_config.py @@ -0,0 +1,40 @@ +# 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 + +import pytest + + +def pytest_configure(config): + config.addinivalue_line("markers", "datatest: this test uses binary data") + + +def pytest_addoption(parser): + """Add --datatest commandline option for pytest. + + Makes sure the option is set only once even when running tests of several model packages in one session. + """ + try: + parser.addoption( + "--datatest", + action="store_true", + help="running tests that use serialized data, can be slow since data might be downloaded from online storage", + default=False, + ) + except ValueError: + pass + + +def pytest_runtest_setup(item): + for _ in item.iter_markers(name="datatest"): + if not item.config.getoption("--datatest"): + pytest.skip("need '--datatest' option to run") diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py new file mode 100644 index 0000000000..f4550ccb58 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -0,0 +1,577 @@ +# 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 +import logging + +import numpy as np +import serialbox as ser +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.atmosphere.diffusion.diffusion import VectorTuple +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.common import dimension +from icon4py.model.common.decomposition.decomposed import DecompositionInfo +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CECDim, + CEDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field + + +class IconSavepoint: + def __init__(self, sp: ser.Savepoint, ser: ser.Serializer, size: dict): + self.savepoint = sp + self.serializer = ser + self.sizes = size + self.log = logging.getLogger((__name__)) + + def log_meta_info(self): + self.log.info(self.savepoint.metainfo) + + def _get_field(self, name, *dimensions, dtype=float): + buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) + buffer_size = ( + self.sizes[d] if d.kind is DimensionKind.HORIZONTAL else s + for s, d in zip(buffer.shape, dimensions) + ) + buffer = buffer[tuple(map(slice, buffer_size))] + + self.log.debug(f"{name} {buffer.shape}") + return np_as_located_field(*dimensions)(buffer) + + def get_metadata(self, *names): + metadata = self.savepoint.metainfo.to_dict() + return {n: metadata[n] for n in names if n in metadata} + + def _read_int32_shift1(self, name: str): + """ + Read a start indices field. + + use for start indices: the shift accounts for the zero based python + values are converted to int32 + """ + return self._read_int32(name, offset=1) + + def _read_int32(self, name: str, offset=0): + """ + Read an end indices field. + + use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based + this accounts for being exclusive python exclusive bounds: [from:to) + field values are convert to int32 + """ + return self._read(name, offset, dtype=int32) + + def _read_bool(self, name: str): + return self._read(name, offset=0, dtype=bool) + + def _read(self, name: str, offset=0, dtype=int): + return (self.serializer.read(name, self.savepoint) - offset).astype(dtype) + + +class IconGridSavePoint(IconSavepoint): + def vct_a(self): + return self._get_field("vct_a", KDim) + + def tangent_orientation(self): + return self._get_field("tangent_orientation", EdgeDim) + + def inverse_primal_edge_lengths(self): + return self._get_field("inv_primal_edge_length", EdgeDim) + + def inv_vert_vert_length(self): + return self._get_field("inv_vert_vert_length", EdgeDim) + + def primal_normal_vert_x(self): + return self._get_field("primal_normal_vert_x", VertexDim, E2C2VDim) + + def primal_normal_vert_y(self): + return self._get_field("primal_normal_vert_y", VertexDim, E2C2VDim) + + def dual_normal_vert_y(self): + return self._get_field("dual_normal_vert_y", VertexDim, E2C2VDim) + + def dual_normal_vert_x(self): + return self._get_field("dual_normal_vert_x", VertexDim, E2C2VDim) + + def cell_areas(self): + return self._get_field("cell_areas", CellDim) + + def edge_areas(self): + return self._get_field("edge_areas", EdgeDim) + + def inv_dual_edge_length(self): + return self._get_field("inv_dual_edge_length", EdgeDim) + + def edge_cell_length(self): + return self._get_field("edge_cell_length", EdgeDim, E2CDim) + + def cells_start_index(self): + return self._read_int32_shift1("c_start_index") + + def cells_end_index(self): + return self._read_int32("c_end_index") + + def vertex_start_index(self): + return self._read_int32_shift1("v_start_index") + + def vertex_end_index(self): + return self._read_int32("v_end_index") + + def edge_start_index(self): + return self._read_int32_shift1("e_start_index") + + def edge_end_index(self): + # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being + # one off accounts for being exclusive [from:to) + return self.serializer.read("e_end_index", self.savepoint) + + def c_owner_mask(self): + return self._get_field("c_owner_mask", CellDim, dtype=bool) + + def e_owner_mask(self): + return self._get_field("e_owner_mask", EdgeDim, dtype=bool) + + def print_connectivity_info(self, name: str, ar: np.ndarray): + self.log.debug(f" connectivity {name} {ar.shape}") + + def c2e(self): + return self._get_connectivity_array("c2e", CellDim) + + def _get_connectivity_array(self, name: str, target_dim: Dimension): + connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :] + self.log.debug(f" connectivity {name} : {connectivity.shape}") + return connectivity + + def c2e2c(self): + return self._get_connectivity_array("c2e2c", CellDim) + + def e2c(self): + return self._get_connectivity_array("e2c", EdgeDim) + + def e2v(self): + # array "e2v" is actually e2c2v + v_ = self._get_connectivity_array("e2v", EdgeDim)[:, 0:2] + self.log.debug(f"real e2v {v_.shape}") + return v_ + + def e2c2v(self): + # array "e2v" is actually e2c2v, that is hexagon or pentagon + return self._get_connectivity_array("e2v", EdgeDim) + + def v2e(self): + return self._get_connectivity_array("v2e", VertexDim) + + def v2c(self): + return self._get_connectivity_array("v2c", VertexDim) + + def c2v(self): + return self._get_connectivity_array("c2v", CellDim) + + def nrdmax(self): + return self._read_int32_shift1("nrdmax") + + def refin_ctrl(self, dim: Dimension): + field_name = "refin_ctl" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def num(self, dim: Dimension): + return self.sizes[dim] + + def _read_field_for_dim(self, field_name, read_func, dim: Dimension): + match (dim): + case dimension.CellDim: + return read_func(f"c_{field_name}") + case dimension.EdgeDim: + return read_func(f"e_{field_name}") + case dimension.VertexDim: + return read_func(f"v_{field_name}") + case _: + raise NotImplementedError( + f"only {dimension.CellDim, dimension.EdgeDim, dimension.VertexDim} are handled" + ) + + def owner_mask(self, dim: Dimension): + field_name = "owner_mask" + mask = self._read_field_for_dim(field_name, self._read_bool, dim) + return np.squeeze(mask) + + def global_index(self, dim: Dimension): + field_name = "glb_index" + return self._read_field_for_dim(field_name, self._read_int32_shift1, dim) + + def decomp_domain(self, dim): + field_name = "decomp_domain" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def construct_decomposition_info(self): + return ( + DecompositionInfo(klevels=self.num(KDim)) + .with_dimension(*self._get_decomp_fields(CellDim)) + .with_dimension(*self._get_decomp_fields(EdgeDim)) + .with_dimension(*self._get_decomp_fields(VertexDim)) + ) + + def _get_decomp_fields(self, dim: Dimension): + global_index = self.global_index(dim) + mask = self.owner_mask(dim)[0 : self.num(dim)] + return dim, global_index, mask + + def construct_icon_grid(self) -> IconGrid: + + cell_starts = self.cells_start_index() + cell_ends = self.cells_end_index() + vertex_starts = self.vertex_start_index() + vertex_ends = self.vertex_end_index() + edge_starts = self.edge_start_index() + edge_ends = self.edge_end_index() + config = GridConfig( + horizontal_config=HorizontalGridSize( + num_vertices=self.num(VertexDim), + num_cells=self.num(CellDim), + num_edges=self.num(EdgeDim), + ), + vertical_config=VerticalGridSize(num_lev=self.num(KDim)), + ) + c2e2c = self.c2e2c() + c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c)) + grid = ( + IconGrid() + .with_config(config) + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) + .with_start_end_indices(EdgeDim, edge_starts, edge_ends) + .with_start_end_indices(CellDim, cell_starts, cell_ends) + .with_connectivities( + { + C2EDim: self.c2e(), + E2CDim: self.e2c(), + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + } + ) + .with_connectivities({E2VDim: self.e2v(), V2EDim: self.v2e(), E2C2VDim: self.e2c2v()}) + ) + return grid + + def construct_edge_geometry(self) -> EdgeParams: + primal_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.primal_normal_vert_y(), ECVDim), + ) + dual_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.dual_normal_vert_y(), ECVDim), + ) + return EdgeParams( + tangent_orientation=self.tangent_orientation(), + inverse_primal_edge_lengths=self.inverse_primal_edge_lengths(), + inverse_dual_edge_lengths=self.inv_dual_edge_length(), + inverse_vertex_vertex_lengths=self.inv_vert_vert_length(), + primal_normal_vert_x=primal_normal_vert[0], + primal_normal_vert_y=primal_normal_vert[1], + dual_normal_vert_x=dual_normal_vert[0], + dual_normal_vert_y=dual_normal_vert[1], + edge_areas=self.edge_areas(), + ) + + def construct_cell_geometry(self) -> CellParams: + return CellParams(area=self.cell_areas()) + + +class InterpolationSavepoint(IconSavepoint): + def geofac_grg(self): + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + num_cells = self.sizes[CellDim] + return np_as_located_field(CellDim, C2E2CODim)(grg[:num_cells, :, 0]), np_as_located_field( + CellDim, C2E2CODim + )(grg[:num_cells, :, 1]) + + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + + def e_bln_c_s(self): + return self._get_field("e_bln_c_s", CellDim, C2EDim) + + def geofac_div(self): + return self._get_field("geofac_div", CellDim, C2EDim) + + def geofac_n2s(self): + return self._get_field("geofac_n2s", CellDim, C2E2CODim) + + def rbf_vec_coeff_v1(self): + return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) + + def rbf_vec_coeff_v2(self): + return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + + def nudgecoeff_e(self): + return self._get_field("nudgecoeff_e", EdgeDim) + + def construct_interpolation_state_for_diffusion( + self, + ) -> DiffusionInterpolationState: + grg = self.geofac_grg() + return DiffusionInterpolationState( + e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), + rbf_coeff_1=self.rbf_vec_coeff_v1(), + rbf_coeff_2=self.rbf_vec_coeff_v2(), + geofac_div=as_1D_sparse_field(self.geofac_div(), CEDim), + geofac_n2s=self.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=self.nudgecoeff_e(), + ) + + def c_lin_e(self): + return self._get_field("c_lin_e", EdgeDim, E2CDim) + + +class MetricSavepoint(IconSavepoint): + def construct_metric_state_for_diffusion(self) -> DiffusionMetricState: + return DiffusionMetricState( + mask_hdiff=self.mask_diff(), + theta_ref_mc=self.theta_ref_mc(), + wgtfac_c=self.wgtfac_c(), + zd_intcoef=self.zd_intcoef(), + zd_vertoffset=self.zd_vertoffset(), + zd_diffcoef=self.zd_diffcoef(), + ) + + def zd_diffcoef(self): + return self._get_field("zd_diffcoef", CellDim, KDim) + + def zd_intcoef(self): + return self._read_and_reorder_sparse_field("vcoef", CellDim) + + def _read_and_reorder_sparse_field(self, name: str, horizontal_dim: Dimension, sparse_size=3): + ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[ + : self.sizes[horizontal_dim], :, : + ] + if ser_input.shape[1] != sparse_size: + ser_input = np.moveaxis((ser_input), 1, -1) + + return self._linearize_first_2dims( + ser_input, sparse_size=sparse_size, target_dims=(CECDim, KDim) + ) + + def _linearize_first_2dims( + self, data: np.ndarray, sparse_size: int, target_dims: tuple[Dimension, ...] + ): + old_shape = data.shape + assert old_shape[1] == sparse_size + return np_as_located_field(*target_dims)( + data.reshape(old_shape[0] * old_shape[1], old_shape[2]) + ) + + def zd_vertoffset(self): + return self._read_and_reorder_sparse_field("zd_vertoffset", CellDim) + + def zd_vertidx(self): + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + + def zd_indlist(self): + return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) + + def theta_ref_mc(self): + return self._get_field("theta_ref_mc", CellDim, KDim) + + def wgtfac_c(self): + return self._get_field("wgtfac_c", CellDim, KDim) + + def wgtfac_e(self): + return self._get_field("wgtfac_e", EdgeDim, KDim) + + def mask_diff(self): + return self._get_field("mask_hdiff", CellDim, KDim, dtype=bool) + + +class IconDiffusionInitSavepoint(IconSavepoint): + def hdef_ic(self): + return self._get_field("hdef_ic", CellDim, KDim) + + def div_ic(self): + return self._get_field("div_ic", CellDim, KDim) + + def dwdx(self): + return self._get_field("dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("dwdy", CellDim, KDim) + + def vn(self): + return self._get_field("vn", EdgeDim, KDim) + + def theta_v(self): + return self._get_field("theta_v", CellDim, KDim) + + def w(self): + return self._get_field("w", CellDim, KDim) + + def exner(self): + return self._get_field("exner", CellDim, KDim) + + def diff_multfac_smag(self): + return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) + + def smag_limit(self): + return np.squeeze(self.serializer.read("smag_limit", self.savepoint)) + + def diff_multfac_n2w(self): + return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) + + def nudgezone_diff(self) -> int: + return self.serializer.read("nudgezone_diff", self.savepoint)[0] + + def bdy_diff(self) -> int: + return self.serializer.read("bdy_diff", self.savepoint)[0] + + def fac_bdydiff_v(self) -> int: + return self.serializer.read("fac_bdydiff_v", self.savepoint)[0] + + def smag_offset(self): + return self.serializer.read("smag_offset", self.savepoint)[0] + + def diff_multfac_w(self): + return self.serializer.read("diff_multfac_w", self.savepoint)[0] + + def diff_multfac_vn(self): + return self.serializer.read("diff_multfac_vn", self.savepoint) + + def construct_prognostics(self) -> PrognosticState: + return PrognosticState( + w=self.w(), + vn=self.vn(), + exner_pressure=self.exner(), + theta_v=self.theta_v(), + ) + + def construct_diagnostics_for_diffusion(self) -> DiffusionDiagnosticState: + return DiffusionDiagnosticState( + hdef_ic=self.hdef_ic(), + div_ic=self.div_ic(), + dwdx=self.dwdx(), + dwdy=self.dwdy(), + ) + + +class IconDiffusionExitSavepoint(IconSavepoint): + def vn(self): + return self._get_field("x_vn", EdgeDim, KDim) + + def theta_v(self): + return self._get_field("x_theta_v", CellDim, KDim) + + def w(self): + return self._get_field("x_w", CellDim, KDim) + + def dwdx(self): + return self._get_field("x_dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("x_dwdy", CellDim, KDim) + + def exner(self): + return self._get_field("x_exner", CellDim, KDim) + + def z_temp(self): + return self._get_field("x_z_temp", CellDim, KDim) + + def div_ic(self): + return self._get_field("x_div_ic", CellDim, KDim) + + def hdef_ic(self): + return self._get_field("x_hdef_ic", CellDim, KDim) + + +class IconSerialDataProvider: + def __init__(self, fname_prefix, path=".", do_print=False, mpi_rank=0): + self.rank = mpi_rank + self.serializer: ser.Serializer = None + self.file_path: str = path + self.fname = f"{fname_prefix}_rank{str(self.rank)}" + self.log = logging.getLogger(__name__) + self._init_serializer(do_print) + self.grid_size = self._grid_size() + + def _init_serializer(self, do_print: bool): + if not self.fname: + self.log.warning(" WARNING: no filename! closing serializer") + self.serializer = ser.Serializer(ser.OpenModeKind.Read, self.file_path, self.fname) + if do_print: + self.print_info() + + def print_info(self): + self.log.info(f"SAVEPOINTS: {self.serializer.savepoint_list()}") + self.log.info(f"FIELDNAMES: {self.serializer.fieldnames()}") + + def _grid_size(self): + sp = self._get_icon_grid_savepoint() + grid_sizes = { + CellDim: self.serializer.read("num_cells", savepoint=sp).astype(int32)[0], + EdgeDim: self.serializer.read("num_edges", savepoint=sp).astype(int32)[0], + VertexDim: self.serializer.read("num_vert", savepoint=sp).astype(int32)[0], + KDim: sp.metainfo.to_dict()["nlev"], + } + return grid_sizes + + def from_savepoint_grid(self) -> IconGridSavePoint: + savepoint = self._get_icon_grid_savepoint() + return IconGridSavePoint(savepoint, self.serializer, size=self.grid_size) + + def _get_icon_grid_savepoint(self): + savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() + return savepoint + + def from_savepoint_diffusion_init( + self, + linit: bool, + date: str, + ) -> IconDiffusionInitSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-init"].linit[linit].date[date].as_savepoint() + ) + return IconDiffusionInitSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_interpolation_savepoint(self) -> InterpolationSavepoint: + savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() + return InterpolationSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_metrics_savepoint(self) -> MetricSavepoint: + savepoint = self.serializer.savepoint["metric_state"].as_savepoint() + return MetricSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_savepoint_diffusion_exit(self, linit: bool, date: str) -> IconDiffusionExitSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-exit"].linit[linit].date[date].as_savepoint() + ) + return IconDiffusionExitSavepoint(savepoint, self.serializer, size=self.grid_size) diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index 5b643bd637..3f3f1a60fe 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -61,6 +61,29 @@ @dataclass class SimpleMeshData: + c2v_table = np.asarray( + [ + [0, 1, 4], + [1, 2, 5], + [2, 0, 3], + [0, 3, 4], + [1, 4, 5], + [2, 5, 3], + [3, 4, 7], + [4, 5, 8], + [5, 3, 6], + [3, 6, 7], + [4, 7, 8], + [5, 8, 6], + [6, 7, 1], + [7, 8, 2], + [8, 6, 0], + [6, 0, 1], + [7, 1, 2], + [8, 2, 0], + ] + ) + e2c2v_table = np.asarray( [ [0, 1, 4, 6], # 0 @@ -379,6 +402,7 @@ class SimpleMesh: def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.diamond_arr = SimpleMeshData.diamond_table + self.c2v = SimpleMeshData.c2v_table self.e2c = SimpleMeshData.e2c_table self.e2v = SimpleMeshData.e2v_table self.c2e = SimpleMeshData.c2e_table @@ -399,6 +423,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_e2c2e = self.e2c2e.shape[1] self.n_e2c2v = self.e2c2v.shape[1] self.n_v2c = self.v2c.shape[1] + self.n_c2v = self.c2v.shape[1] self.n_v2e = self.v2e.shape[1] self.n_cells = self.c2e.shape[0] self.n_c2e2c2e2c = self.c2e2c2e2c.shape[1] @@ -425,6 +450,9 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): C2E2C2E2CDim: self.n_c2e2c2e2c, } + def get_c2v_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2v, VertexDim, CellDim, self.n_c2v) + def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e, CellDim, EdgeDim, self.n_c2e) @@ -457,6 +485,15 @@ def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_c2e2c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e2c2e2c, CellDim, CellDim, self.n_c2e2c2e2c) + def get_e2ecv_offset_provider(self): + old_shape = self.e2c2v.shape + e2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(e2ecv_table, EdgeDim, ECVDim, e2ecv_table.shape[1]) + + def get_c2ce_offset_provider(self): + old_shape = self.c2e.shape + c2ce_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(c2ce_table, CellDim, CEDim, c2ce_table.shape[1]) def get_offset_provider(self): return { @@ -470,9 +507,9 @@ def get_offset_provider(self): "E2C": self.get_e2c_offset_provider(), "E2V": self.get_e2v_offset_provider(), "E2C2V": self.get_e2c2v_offset_provider(), + "C2CE": self.get_c2ce_offset_provider(), "Koff": KDim, "C2E2C2E2C": self.get_c2e2c2e2c_offset_provider(), - "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, self.n_c2e), "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, self.n_e2c2v), "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, self.n_e2c), } diff --git a/model/common/src/icon4py/model/common/utils.py b/model/common/src/icon4py/model/common/utils.py new file mode 100644 index 0000000000..403bf7a685 --- /dev/null +++ b/model/common/src/icon4py/model/common/utils.py @@ -0,0 +1,22 @@ +# 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 + + +def builder(func): + """Use as decorator on builder functions.""" + + def wrapper(self, *args, **kwargs): + func(self, *args, **kwargs) + return self + + return wrapper diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py new file mode 100644 index 0000000000..63571d6840 --- /dev/null +++ b/model/common/tests/conftest.py @@ -0,0 +1,64 @@ +# 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 +import pytest + +from icon4py.model.common.test_utils.data_handling import download_and_extract +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 + backend, + base_path, + damping_height, + data_provider, + datapath, + decomposition_info, + download_ser_data, + grid_savepoint, + icon_grid, + interpolation_savepoint, + mesh, + processor_props, + ranked_data_path, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) + + +grids_path = base_path.joinpath("grids") +r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name +r02b04_global_grid_path = grids_path.joinpath("r02b04_global") +r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name + + +mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" + + +@pytest.fixture() +def r04b09_dsl_gridfile(get_grid_files): + return r04b09_dsl_grid_path.joinpath("grid.nc") + + +@pytest.fixture(scope="session") +def get_grid_files(pytestconfig): + """ + Get the grid files used for testing. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + if not pytestconfig.getoption("datatest"): + pytest.skip("not running datatest marked tests") + download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) + download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py new file mode 100644 index 0000000000..b363461f0f --- /dev/null +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -0,0 +1,196 @@ +# 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 + +import numpy as np +import pytest + +from icon4py.model.common.decomposition.decomposed import ( + DecompositionInfo, + DomainDescriptorIdGenerator, + GHexMultiNode, + SingleNode, + create_exchange, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size + + +""" +running tests with mpi: + +mpirun -np 2 python -m pytest -v --with-mpi tests/mpi_tests/test_parallel_setup.py + +mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ + + +""" + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_props(processor_props): + assert processor_props.comm + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + ("dim, owned, total"), + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), +) +def test_decomposition_info_masked( + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 +): + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) + my_total = total[my_rank] + my_owned = owned[my_rank] + assert all_indices.shape[0] == my_total + + owned_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.OWNED) + assert owned_indices.shape[0] == my_owned + + halo_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == my_total - my_owned + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + + +# @pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +def _assert_index_partitioning(all_indices, halo_indices, owned_indices): + owned_list = owned_indices.tolist() + halos_list = halo_indices.tolist() + all_list = all_indices.tolist() + assert set(owned_list) & set(halos_list) == set() + assert set(owned_list) & set(all_list) == set(owned_list) + assert set(halos_list) & set(all_list) == set(halos_list) + assert set(halos_list) | set(owned_list) == set(all_list) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + ("dim, owned, total"), + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), +) +@pytest.mark.mpi(min_size=2) +def test_decomposition_info_local_index( + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 +): + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) + my_total = total[my_rank] + my_owned = owned[my_rank] + + assert all_indices.shape[0] == my_total + assert np.array_equal(all_indices, np.arange(0, my_total)) + halo_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == my_total - my_owned + assert halo_indices.shape[0] < all_indices.shape[0] + assert np.alltrue(halo_indices <= np.max(all_indices)) + + owned_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.OWNED) + assert owned_indices.shape[0] == my_owned + assert owned_indices.shape[0] <= all_indices.shape[0] + assert np.alltrue(owned_indices <= np.max(all_indices)) + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("num", [1, 2, 3]) +def test_domain_descriptor_id_are_globally_unique(num, processor_props): + props = processor_props + size = props.comm_size + id_gen = DomainDescriptorIdGenerator(parallel_props=props) + id1 = id_gen() + assert id1 == props.comm_size * props.rank + assert id1 < props.comm_size * (props.rank + 1) + ids = [] + ids.append(id1) + for _ in range(1, num * size): + next_id = id_gen() + assert next_id > id1 + ids.append(next_id) + all_ids = props.comm.gather(ids, root=0) + if props.rank == 0: + all_ids = np.asarray(all_ids).flatten() + assert len(all_ids) == size * size * num + assert len(all_ids) == len(set(all_ids)) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_decomposition_info_matches_gridsize( + caplog, + download_ser_data, + decomposition_info, + icon_grid, + processor_props, +): # F811 + + check_comm_size(processor_props) + assert ( + decomposition_info.global_index( + dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL + ).shape[0] + == icon_grid.num_cells() + ) + assert ( + decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[0] + == icon_grid.num_vertices() + ) + assert ( + decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[0] + == icon_grid.num_edges() + ) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_create_multi_node_runtime_with_mpi( + download_ser_data, decomposition_info, processor_props +): # F811 + props = processor_props + exchange = create_exchange(props, decomposition_info) + if props.comm_size > 1: + assert isinstance(exchange, GHexMultiNode) + else: + assert isinstance(exchange, SingleNode) + + +@pytest.mark.parametrize("processor_props", [False], indirect=True) +def test_create_single_node_runtime_without_mpi(processor_props, decomposition_info): + props = processor_props + exchange = create_exchange(props, decomposition_info) + + assert isinstance(exchange, SingleNode) diff --git a/model/common/tests/mpi_tests/test_parallel_setup.py b/model/common/tests/mpi_tests/test_parallel_setup.py new file mode 100644 index 0000000000..bcf714005d --- /dev/null +++ b/model/common/tests/mpi_tests/test_parallel_setup.py @@ -0,0 +1,47 @@ +# 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 + + +import pytest +from mpi4py import MPI + +from icon4py.model.common.decomposition.parallel_setup import get_processor_properties, init_mpi + + +@pytest.mark.mpi +def test_parallel_properties_from_comm_world(): + props = get_processor_properties(with_mpi=True) + assert props.rank < props.comm_size + assert props.comm_name == "MPI_COMM_WORLD" + + +@pytest.mark.mpi(min_size=2) +def test_parallel_properties_from_mpi_comm(): + init_mpi() + world = MPI.COMM_WORLD + group = world.Get_group() + pair = group.Incl([0, 1]) + comm = world.Create(pair) + if comm != MPI.COMM_NULL: + comm.Set_name("my_comm") + props = get_processor_properties(with_mpi=True, comm_id=comm) + assert props.rank < props.comm_size + assert props.comm_size == 2 + assert props.comm_name == "my_comm" + + +def test_single_node_properties(): + props = get_processor_properties() + assert props.comm_size == 1 + assert props.rank == 0 + assert props.comm_name == "" diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py new file mode 100644 index 0000000000..6f5f6c47ae --- /dev/null +++ b/model/common/tests/test_grid_manager.py @@ -0,0 +1,521 @@ +# 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 +import logging +from uuid import uuid4 + +import numpy as np +import pytest +from netCDF4 import Dataset + +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.grid_manager import ( + GridFile, + GridFileName, + GridManager, + IndexTransformation, + ToGt4PyTransformation, +) +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.grid.icon_grid import VerticalGridSize +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +SIMPLE_MESH_NC = "simple_mesh_grid.nc" + + +@pytest.fixture +def simple_mesh_gridfile(tmp_path): + path = tmp_path.joinpath(SIMPLE_MESH_NC).absolute() + mesh = SimpleMesh() + dataset = Dataset(path, "w", format="NETCDF4") + dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) + dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=mesh.n_vertices) + + dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=mesh.n_edges) + dataset.createDimension(GridFile.DimensionName.CELL_NAME, size=mesh.n_cells) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, size=mesh.n_e2v) + dataset.createDimension(GridFile.DimensionName.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) + dataset.createDimension(GridFile.DimensionName.MAX_CHILD_DOMAINS, size=1) + # add dummy values for the grf dimensions + dataset.createDimension(GridFile.DimensionName.CELL_GRF, size=14) + dataset.createDimension(GridFile.DimensionName.EDGE_GRF, size=24) + dataset.createDimension(GridFile.DimensionName.VERTEX_GRF, size=13) + _add_to_dataset( + dataset, + np.zeros(mesh.n_edges), + GridFile.GridRefinementName.CONTROL_EDGES, + (GridFile.DimensionName.EDGE_NAME,), + ) + + _add_to_dataset( + dataset, + np.zeros(mesh.n_cells), + GridFile.GridRefinementName.CONTROL_CELLS, + (GridFile.DimensionName.CELL_NAME,), + ) + _add_to_dataset( + dataset, + np.zeros(mesh.n_vertices), + GridFile.GridRefinementName.CONTROL_VERTICES, + (GridFile.DimensionName.VERTEX_NAME,), + ) + + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, size=mesh.n_c2e) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, size=mesh.n_v2c) + + _add_to_dataset( + dataset, + mesh.c2e, + GridFile.OffsetName.C2E, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.e2c, + GridFile.OffsetName.E2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), + ) + _add_to_dataset( + dataset, + mesh.e2v, + GridFile.OffsetName.E2V, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.v2c, + GridFile.OffsetName.V2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.c2v, + GridFile.OffsetName.C2V, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + _add_to_dataset( + dataset, + np.zeros((mesh.n_vertices, 4), dtype=np.int32), + GridFile.OffsetName.V2E2V, + (GridFile.DimensionName.DIAMOND_EDGE_SIZE, GridFile.DimensionName.VERTEX_NAME), + ) + _add_to_dataset( + dataset, + mesh.v2e, + GridFile.OffsetName.V2E, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), + ) + _add_to_dataset( + dataset, + mesh.c2e2c, + GridFile.OffsetName.C2E2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) + dataset.close() + yield path + path.unlink() + + +def _add_to_dataset( + dataset: Dataset, + data: np.ndarray, + var_name: str, + dims: tuple[GridFileName, GridFileName], +): + var = dataset.createVariable(var_name, np.int32, dims) + var[:] = np.transpose(data)[:] + + +def test_gridparser_dimension(simple_mesh_gridfile): + data = Dataset(simple_mesh_gridfile, "r") + grid_parser = GridFile(data) + mesh = SimpleMesh() + assert grid_parser.dimension(GridFile.DimensionName.CELL_NAME) == mesh.n_cells + assert grid_parser.dimension(GridFile.DimensionName.VERTEX_NAME) == mesh.n_vertices + assert grid_parser.dimension(GridFile.DimensionName.EDGE_NAME) == mesh.n_edges + + +@pytest.mark.datatest +def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): + data = Dataset(r04b09_dsl_gridfile, "r") + grid_file = GridFile(data) + + assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num(CellDim) + assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num(EdgeDim) + assert grid_file.dimension(GridFile.DimensionName.VERTEX_NAME) == grid_savepoint.num(VertexDim) + + +def test_grid_parser_index_fields(simple_mesh_gridfile, caplog): + caplog.set_level(logging.DEBUG) + data = Dataset(simple_mesh_gridfile, "r") + mesh = SimpleMesh() + grid_parser = GridFile(data) + + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.C2E), mesh.c2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.E2C), mesh.e2c) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2E), mesh.v2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2C), mesh.v2c) + + +# TODO @magdalena add test cases for hexagon vertices v2e2v +# v2e2v: grid,??? + +# v2e: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + seralized_v2e = grid_savepoint.v2e()[0 : grid.num_vertices(), :] + # there are vertices at the boundary of a local domain or at a pentagon point that have less than + # 6 neighbors hence there are "Missing values" in the grid file + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(seralized_v2e) + assert has_invalid_index(grid.get_v2e_connectivity().table) + reset_invalid_index(seralized_v2e) + assert np.allclose(grid.get_v2e_connectivity().table, seralized_v2e) + + +# v2c: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_v2c = grid_savepoint.v2c()[0 : grid.num_vertices(), :] + # there are vertices that have less than 6 neighboring cells: either pentagon points or + # vertices at the boundary of the domain for a limited area mode + # hence in the grid file there are "missing values" + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(serialized_v2c) + assert has_invalid_index(grid.get_v2c_connectivity().table) + reset_invalid_index(serialized_v2c) + + assert np.allclose(grid.get_v2c_connectivity().table, serialized_v2c) + + +def reset_invalid_index(index_array: np.ndarray): + """ + Revert changes from mo_model_domimp_patches. + + Helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. + see: + # ! Checks for the pentagon case and moves dummy cells to end. + # ! The dummy entry is either set to 0 or duplicated from the last one + # SUBROUTINE move_dummies_to_end(array, array_size, max_connectivity, duplicate) + + After reading the grid file ICON moves all invalid indices (neighbors not existing in the + grid file) to the end of the neighbor list and replaces them with the "last valid neighbor index" + it is up to the user then to ensure that any coefficients in neighbor some multiplied with + these values are zero in order to "remove" them again from the sum. + + For testing we resubstitute those to the GridFile.INVALID_INDEX + Args: + index_array: the array where values the invalid values have to be reset + + Returns: an array where the spurious "last valid index" are replaced by GridFile.INVALID_INDEX + + """ + for i in range(0, index_array.shape[0]): + uq, index = np.unique(index_array[i, :], return_index=True) + index_array[i, max(index) + 1 :] = GridFile.INVALID_INDEX + + +# e2v: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + + serialized_e2v = grid_savepoint.e2v()[0 : grid.num_edges(), :] + # all vertices in the system have to neighboring edges, there no edges that point nowhere + # hence this connectivity has no "missing values" in the grid file + assert not has_invalid_index(serialized_e2v) + assert not has_invalid_index(grid.get_e2v_connectivity().table) + assert np.allclose(grid.get_e2v_connectivity().table, serialized_e2v) + + +def has_invalid_index(ar: np.ndarray): + return np.any(np.where(ar == GridFile.INVALID_INDEX)) + + +# e2c : exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_e2c = grid_savepoint.e2c()[0 : grid.num_edges(), :] + # there are edges at the boundary that have only one + # neighboring cell, there are "missing values" in the grid file + # and here they do not get substituted in the ICON preprocessing + assert has_invalid_index(serialized_e2c) + assert has_invalid_index(grid.get_e2c_connectivity().table) + assert np.allclose(grid.get_e2c_connectivity().table, serialized_e2c) + + +# c2e: serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + + serialized_c2e = grid_savepoint.c2e()[0 : grid.num_cells(), :] + # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the + # first place + # hence there are no "missing values" in the grid file + assert not has_invalid_index(serialized_c2e) + assert not has_invalid_index(grid.get_c2e_connectivity().table) + assert np.allclose(grid.get_c2e_connectivity().table, serialized_c2e) + + +# e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh +@pytest.mark.datatest +@pytest.mark.skip("does not directly exist in the grid file, needs to be constructed") +# TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell +def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile) + serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :] + assert has_invalid_index(serialized_e2c2e) + grid = gm.get_grid() + assert has_invalid_index(grid.get_e2c2e_connectivity().table) + assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) + + +# c2e2c: exists in serial, simple_mesh, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert np.allclose( + grid.get_c2e2c_connectivity().table, + grid_savepoint.c2e2c()[0 : grid.num_cells(), :], + ) + + +@pytest.mark.xfail +@pytest.mark.datatest +def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + # the "far" (adjacent to edge normal ) is not there. why? + # despite that: ordering is different + assert np.allclose( + grid.get_e2c2v_connectivity().table, + grid_savepoint.e2c2v()[0 : grid.num_edges(), :], + ) + + +@pytest.mark.datatest +def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + c2v = grid.get_c2v_connectivity().table + assert np.allclose(c2v, grid_savepoint.c2v()[0 : grid.num_cells(), :]) + + +def init_grid_manager(fname): + gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(65)) + gm() + return gm + + +@pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) +def test_grid_manager_getsize(simple_mesh_gridfile, dim, size, caplog): + caplog.set_level(logging.DEBUG) + gm = GridManager(IndexTransformation(), simple_mesh_gridfile, VerticalGridSize(num_lev=80)) + gm() + assert size == gm.get_size(dim) + + +def test_grid_manager_diamond_offset(simple_mesh_gridfile): + mesh = SimpleMesh() + gm = GridManager( + IndexTransformation(), + simple_mesh_gridfile, + VerticalGridSize(num_lev=mesh.k_level), + ) + gm() + grid = gm.get_grid() + assert np.allclose( + np.sort(grid.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) + ) + + +def test_gridmanager_given_file_not_found_then_abort(): + fname = "./unknown_grid.nc" + with pytest.raises(SystemExit) as error: + gm = GridManager(IndexTransformation(), fname, VerticalGridSize(num_lev=80)) + gm() + assert error.type == SystemExit + assert error.value == 1 + + +@pytest.mark.parametrize("size", [100, 1500, 20000]) +def test_gt4py_transform_offset_by_1_where_valid(size): + trafo = ToGt4PyTransformation() + input_field = np.random.randint(-1, size, (size,)) + offset = trafo.get_offset_for_index_field(input_field) + expected = np.where(input_field >= 0, -1, 0) + assert np.allclose(expected, offset) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), + ], +) +def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_start_index(dim, marker) == index + assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), + ], +) +def test_get_end_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_end_index(dim, marker) == index + assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker) diff --git a/model/common/tests/test_icon_grid.py b/model/common/tests/test_icon_grid.py new file mode 100644 index 0000000000..83f20ee41e --- /dev/null +++ b/model/common/tests/test_icon_grid.py @@ -0,0 +1,400 @@ +# 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 +import pytest + +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex + + +@pytest.mark.datatest +# TODO(Magdalena) HorizontalMarkerIndex.local(dim) does not yield equivalent results form grid file +# and serialized data, why?. Serialized data has those strange -1 values +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), + ], +) +def test_horizontal_end_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_end_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), -1), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), -1), # ???? + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), -1), # ??? + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), + ], +) +def test_horizontal_start_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_start_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(CellDim), + HorizontalMarkerIndex.lateral_boundary(CellDim), + (0, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + (850, 1688), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + (1688, 2511), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + (2511, 3316), + ), + ( + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.interior(CellDim), + (4104, 20896), + ), + ( + HorizontalMarkerIndex.interior(CellDim) + 1, + HorizontalMarkerIndex.interior(CellDim) + 1, + (0, 850), + ), + ( + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.nudging(CellDim), + ( + 3316, + 4104, + ), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.local(CellDim), + (-1, 20896), + ), + ], +) +def test_horizontal_cell_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + CellDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + (428, 1278), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + (1278, 1700), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + (1700, 2538), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + (2538, 2954), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + (2954, 3777), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + (3777, 4184), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + (4184, 4989), + ), + ( + HorizontalMarkerIndex.interior(EdgeDim), + HorizontalMarkerIndex.interior(EdgeDim), + (6176, 31558), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ( + 4989, + 5387, + ), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + (5387, 6176), + ), + ( + HorizontalMarkerIndex.end(EdgeDim), + HorizontalMarkerIndex.end(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.local(EdgeDim), + HorizontalMarkerIndex.local(EdgeDim), + (-1, 31558), + ), + ], +) +def test_horizontal_edge_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + EdgeDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + (428, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + (850, 1266), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + (1266, 1673), + ), + ( + HorizontalMarkerIndex.interior(VertexDim), + HorizontalMarkerIndex.interior(VertexDim), + (2071, 10663), + ), + ( + HorizontalMarkerIndex.interior(VertexDim) + 1, + HorizontalMarkerIndex.interior(VertexDim) + 1, + (0, 428), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.halo(VertexDim), + HorizontalMarkerIndex.halo(VertexDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.local(VertexDim), + HorizontalMarkerIndex.local(VertexDim), + (-1, 10663), + ), + ], +) +def test_horizontal_vertex_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + VertexDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +def test_cross_check_marker_equivalences(icon_grid): + """Check actual equivalences of calculated markers.""" + # TODO(Magdalena): This should go away once we refactor these markers in a good way, such that no calculation need to be done with them anymore. + + assert icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.local(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) - 1, + ) == icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ) + assert icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.nudging(CellDim) - 1, + ) == icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + ) + assert icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 1, + ) == icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) + + assert icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + ) == icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ) + + +@pytest.mark.datatest +def test_grid_size(grid_savepoint): + assert 10663 == grid_savepoint.num(VertexDim) + assert 20896 == grid_savepoint.num(CellDim) + assert 31558 == grid_savepoint.num(EdgeDim) diff --git a/model/common/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py new file mode 100644 index 0000000000..e4c8eb5b97 --- /dev/null +++ b/model/common/tests/test_interpolation_fields.py @@ -0,0 +1,51 @@ +# 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 +# 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 + +import numpy as np +import pytest + +from icon4py.model.common.dimension import EdgeDim +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.interpolation.interpolation_fields import compute_c_lin_e + + +@pytest.mark.datatest +def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() + edge_cell_length = grid_savepoint.edge_cell_length() + owner_mask = grid_savepoint.e_owner_mask() + c_lin_e_ref = interpolation_savepoint.c_lin_e() + lateral_boundary = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + c_lin_e = compute_c_lin_e( + np.asarray(edge_cell_length), + np.asarray(inv_dual_edge_length), + np.asarray(owner_mask), + lateral_boundary, + ) + + assert np.allclose(c_lin_e, c_lin_e_ref) diff --git a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 85% rename from model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 809e48b20f..189beae082 100644 --- a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,11 +13,12 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( +from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -51,4 +52,8 @@ def input_data(self, mesh): ptr_coeff_2=ptr_coeff_2, p_v_out=p_v_out, p_u_out=p_u_out, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_vertices), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/common/tests/test_vertical.py b/model/common/tests/test_vertical.py new file mode 100644 index 0000000000..4d6d71b9a4 --- /dev/null +++ b/model/common/tests/test_vertical.py @@ -0,0 +1,47 @@ +# 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 + +import math + +import numpy as np +import pytest + +from icon4py.model.common.dimension import KDim +from icon4py.model.common.grid.vertical import VerticalModelParams + + +@pytest.mark.parametrize( + "max_h,damping,delta", + [(60000, 34000, 612), (12000, 10000, 100), (109050, 45000, 123)], +) +def test_nrdmax_calculation(max_h, damping, delta): + vct_a = np.arange(0, max_h, delta) + vct_a = vct_a[::-1] + vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) + assert vertical_params.index_of_damping_layer == vct_a.shape[0] - math.ceil(damping / delta) - 1 + + +@pytest.mark.datatest +def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): + a = grid_savepoint.vct_a() + nrdmax = grid_savepoint.nrdmax() + vertical_params = VerticalModelParams(rayleigh_damping_height=damping_height, vct_a=a) + assert nrdmax == vertical_params.index_of_damping_layer + a_array = np.asarray(a) + assert a_array[nrdmax] > damping_height + assert a_array[nrdmax + 1] < damping_height + + +@pytest.mark.datatest +def test_grid_size(grid_savepoint): + assert 65 == grid_savepoint.num(KDim) diff --git a/model/driver/.bumpversion.cfg b/model/driver/.bumpversion.cfg new file mode 100644 index 0000000000..ae24af96da --- /dev/null +++ b/model/driver/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/driver/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/driver/.flake8 b/model/driver/.flake8 new file mode 100644 index 0000000000..31cecff5ab --- /dev/null +++ b/model/driver/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/driver/.pre-commit-config.yaml b/model/driver/.pre-commit-config.yaml new file mode 100644 index 0000000000..4ae0972f5f --- /dev/null +++ b/model/driver/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/driver/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/driver/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/driver/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/driver/.flake8, model/driver/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/driver/README.md b/model/driver/README.md new file mode 100644 index 0000000000..aec461abb0 --- /dev/null +++ b/model/driver/README.md @@ -0,0 +1,36 @@ +# Dummy driver for Python ICON port + +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far the code mostly draws on serialized ICON data until we increasingly can initialize and run the model independently. + +It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. + +Currently, it does _no real timestepping_, instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. + +The code is meant to be changed and enlarged as we port new parts of the model. + +It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes are available. + +## Installation + +See the general instructions in the [README.md](../../README.md) in the base folder of the repository. + +## Usage + +```bash +export ICON4PY_ROOT= +dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +``` + +or if running in parallel + +```bash +mpirun -np 2 dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon +``` + +#### Remarks + +- First (required) arg is the folder where the serialized input data is stored. The input data is the same as is used in the unit tests. The path in the example is where the data is put when downloaded via the unit tests. + - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. +- parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. +- The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. +- The code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml new file mode 100644 index 0000000000..b011e5f5ba --- /dev/null +++ b/model/driver/pyproject.toml @@ -0,0 +1,126 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] + +[project] +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "pytz>=2023.2", + "icon4py-common>=0.0.5", + "icon4py-atmosphere-dycore>=0.0.5", + "icon4py-atmosphere-diffusion>=0.0.5" +] +description = "ICON model driver." +dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-driver" +readme = "README.md" +requires-python = ">=3.10" + +[project.scripts] +dycore_driver = "icon4py.model.driver.dycore_driver:main" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +branch = true +parallel = true +source_pkgs = ['driver'] + +[tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +multi_line_output = 3 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +use_parentheses = true + +[tool.mypy] +disallow_incomplete_defs = true +disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] +ignore_missing_imports = false +implicit_reexport = true +install_types = true +non_interactive = true +show_column_numbers = true +show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.driver.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.driver' = ['py.typed'] diff --git a/model/driver/requirements-dev.txt b/model/driver/requirements-dev.txt new file mode 100644 index 0000000000..68e2b1610d --- /dev/null +++ b/model/driver/requirements-dev.txt @@ -0,0 +1,5 @@ +-r ../../base-requirements-dev.txt +-e ../common +-e ../atmosphere/diffusion +-e ../atmosphere/dycore +-e . diff --git a/model/driver/requirements.txt b/model/driver/requirements.txt new file mode 100644 index 0000000000..18a81dc7c3 --- /dev/null +++ b/model/driver/requirements.txt @@ -0,0 +1,5 @@ +-r ../../base-requirements.txt +../common +../atmosphere/dycore +../atmosphere/diffusion +. diff --git a/model/driver/src/icon4py/model/driver/__init__.py b/model/driver/src/icon4py/model/driver/__init__.py new file mode 100644 index 0000000000..dab7089554 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/__init__.py @@ -0,0 +1,33 @@ +# 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 typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py new file mode 100644 index 0000000000..9103a2443e --- /dev/null +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -0,0 +1,295 @@ +# 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 +import logging +from datetime import datetime, timedelta +from pathlib import Path +from typing import Callable + +import click +import pytz +from devtools import Timer +from gt4py.next import Field, program +from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k +from icon4py.model.common.decomposition.decomposed import create_exchange +from icon4py.model.common.decomposition.parallel_setup import ( + ProcessProperties, + get_processor_properties, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils import serialbox_utils as sb +from icon4py.model.driver.icon_configuration import IconRunConfig, read_config +from icon4py.model.driver.io_utils import ( + SIMULATION_START_DATE, + configure_logging, + read_decomp_info, + read_geometry_fields, + read_icon_grid, + read_initial_state, + read_static_fields, +) + + +log = logging.getLogger(__name__) + + +# TODO (magdalena) to be removed once there is a proper time stepping +@program +def _copy_diagnostic_and_prognostics( + hdef_ic_new: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + div_ic_new: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + dwdx_new: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy_new: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + vn_new: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + w_new: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + exner_new: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + theta_v_new: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], +): + _identity_c_k(hdef_ic_new, out=hdef_ic) + _identity_c_k(div_ic_new, out=div_ic) + _identity_c_k(dwdx_new, out=dwdx) + _identity_c_k(dwdy_new, out=dwdy) + _identity_e_k(vn_new, out=vn) + _identity_c_k(w_new, out=w) + _identity_c_k(exner_new, out=exner) + _identity_c_k(theta_v_new, out=theta_v) + + +class DummyAtmoNonHydro: + def __init__(self, data_provider: sb.IconSerialDataProvider): + self.config = None + self.data_provider = data_provider + self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) + + def init(self, config): + self.config = config + + def _next_physics_date(self, dtime: float): + dynamics_dtime = dtime / self.config.n_substeps + self.simulation_date += timedelta(seconds=dynamics_dtime) + + def _dynamics_timestep(self, dtime): + """Show structure with this dummy fucntion called inside substepping loop.""" + self._next_physics_date(dtime) + + def do_dynamics_substepping( + self, + dtime, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + for _ in range(self.config.n_substeps): + self._dynamics_timestep(dtime) + sp = self.data_provider.from_savepoint_diffusion_init( + linit=False, date=self.simulation_date.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + ) + new_p = sp.construct_prognostics() + new_d = sp.construct_diagnostics_for_diffusion() + _copy_diagnostic_and_prognostics.with_backend(run_gtfn)( + new_d.hdef_ic, + diagnostic_state.hdef_ic, + new_d.div_ic, + diagnostic_state.div_ic, + new_d.dwdx, + diagnostic_state.dwdx, + new_d.dwdy, + diagnostic_state.dwdy, + new_p.vn, + prognostic_state.vn, + new_p.w, + prognostic_state.w, + new_p.exner_pressure, + prognostic_state.exner_pressure, + new_p.theta_v, + prognostic_state.theta_v, + offset_provider={}, + ) + + +class Timeloop: + @classmethod + def name(cls): + return cls.__name__ + + def __init__( + self, + config: IconRunConfig, + diffusion: Diffusion, + atmo_non_hydro: DummyAtmoNonHydro, + ): + self.config = config + self.diffusion = diffusion + self.atmo_non_hydro = atmo_non_hydro + + def _full_name(self, func: Callable): + return ":".join((self.__class__.__name__, func.__name__)) + + def _timestep( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + + self.atmo_non_hydro.do_dynamics_substepping( + self.config.dtime, diagnostic_state, prognostic_state + ) + self.diffusion.run( + diagnostic_state, + prognostic_state, + self.config.dtime, + ) + + def __call__( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + log.info( + f"starting time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + log.info("running initial step to diffuse fields before timeloop starts") + self.diffusion.initial_run( + diagnostic_state, + prognostic_state, + self.config.dtime, + ) + log.info( + f"starting real time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + timer = Timer(self._full_name(self._timestep)) + for t in range(self.config.n_time_steps): + log.info(f"run timestep : {t}") + timer.start() + self._timestep(diagnostic_state, prognostic_state) + timer.capture() + timer.summary(True) + + +def initialize(n_time_steps, file_path: Path, props: ProcessProperties): + """ + Inititalize the driver run. + + "reads" in + - configuration + + - grid information + + - (serialized) input fields, initial + + Returns: + tl: configured timeloop, + prognostic_state: initial state fro prognostic and diagnostic variables + diagnostic_state: + """ + log.info("initialize parallel runtime") + experiment_name = "mch_ch_r04b09_dsl" + log.info(f"reading configuration: experiment {experiment_name}") + config = read_config(experiment_name, n_time_steps=n_time_steps) + + decomp_info = read_decomp_info(file_path, props) + + log.info(f"initializing the grid from '{file_path}'") + icon_grid = read_icon_grid(file_path, rank=props.rank) + log.info(f"reading input fields from '{file_path}'") + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( + file_path, rank=props.rank + ) + (metric_state, interpolation_state) = read_static_fields(file_path) + + log.info("initializing diffusion") + diffusion_params = DiffusionParams(config.diffusion_config) + exchange = create_exchange(props, decomp_info) + diffusion = Diffusion(exchange) + diffusion.init( + icon_grid, + config.diffusion_config, + diffusion_params, + vertical_geometry, + metric_state, + interpolation_state, + edge_geometry, + cell_geometry, + ) + + data_provider, diagnostic_state, prognostic_state = read_initial_state( + file_path, rank=props.rank + ) + + atmo_non_hydro = DummyAtmoNonHydro(data_provider) + atmo_non_hydro.init(config=config.dycore_config) + + tl = Timeloop( + config=config.run_config, + diffusion=diffusion, + atmo_non_hydro=atmo_non_hydro, + ) + return tl, diagnostic_state, prognostic_state + + +@click.command() +@click.argument("input_path") +@click.option("--run_path", default="", help="folder for output") +@click.option("--n_steps", default=5, help="number of time steps to run, max 5 is supported") +@click.option("--mpi", default=False, help="whether or not you are running with mpi") +def main(input_path, run_path, n_steps, mpi): + """ + Run the driver. + + usage: python driver/dycore_driver.py ../../tests/ser_icondata/mch_ch_r04b09_dsl/ser_data + + steps: + 1. initialize model: + + a) load config + + b) initialize grid + + c) initialize/configure components ie "granules" + + d) setup the time loop + + 2. run time loop + + """ + start_time = datetime.now().astimezone(pytz.UTC) + parallel_props = get_processor_properties(with_mpi=mpi) + configure_logging(run_path, start_time, parallel_props) + log.info(f"Starting ICON dycore run: {datetime.isoformat(start_time)}") + log.info(f"input args: input_path={input_path}, n_time_steps={n_steps}") + timeloop, diagnostic_state, prognostic_state = initialize( + n_steps, Path(input_path), parallel_props + ) + log.info("dycore configuring: DONE") + log.info("timeloop: START") + + timeloop(diagnostic_state, prognostic_state) + + log.info("timeloop: DONE") + + +if __name__ == "__main__": + main() diff --git a/model/driver/src/icon4py/model/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py new file mode 100644 index 0000000000..187c4f4f45 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -0,0 +1,85 @@ +# 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 dataclasses import dataclass +from typing import Optional + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType + + +n_substeps_reduced = 2 + + +@dataclass +class IconRunConfig: + n_time_steps: int = 5 + dtime: float = 600.0 + + +@dataclass +class AtmoNonHydroConfig: + n_substeps: int = 5 + + +@dataclass +class IconConfig: + run_config: IconRunConfig + diffusion_config: DiffusionConfig + dycore_config: AtmoNonHydroConfig + + +def read_config(experiment: Optional[str], n_time_steps: int) -> IconConfig: + def _default_run_config(n_steps: int): + if n_steps > 5: + raise NotImplementedError("only five dummy timesteps available") + return IconRunConfig(n_time_steps=n_steps) + + def mch_ch_r04b09_diffusion_config(): + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + n_substeps=n_substeps_reduced, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + ) + + def _default_diffusion_config(): + return DiffusionConfig() + + def _default_config(n_steps): + run_config = _default_run_config(n_steps) + return run_config, _default_diffusion_config(), AtmoNonHydroConfig() + + def _mch_ch_r04b09_config(n_steps): + return ( + IconRunConfig(n_time_steps=n_steps, dtime=10.0), + mch_ch_r04b09_diffusion_config(), + AtmoNonHydroConfig(), + ) + + if experiment == "mch_ch_r04b09_dsl": + (model_run_config, diffusion_config, dycore_config) = _mch_ch_r04b09_config(n_time_steps) + else: + (model_run_config, diffusion_config, dycore_config) = _default_config(n_time_steps) + return IconConfig( + run_config=model_run_config, + diffusion_config=diffusion_config, + dycore_config=dycore_config, + ) diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py new file mode 100644 index 0000000000..ae14933bb3 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -0,0 +1,187 @@ +# 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 + +import logging +from datetime import datetime +from enum import Enum +from pathlib import Path + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.common.decomposition.decomposed import DecompositionInfo +from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.icon_grid import IconGrid +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils import serialbox_utils as sb + + +SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." + +SIMULATION_START_DATE = "2021-06-20T12:00:10.000" +log = logging.getLogger(__name__) + + +class SerializationType(str, Enum): + SB = "serialbox" + NC = "netcdf" + + +def read_icon_grid( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> IconGrid: + """ + Read icon grid. + + Args: + path: path where to find the input data + ser_type: type of input data. Currently only 'sb (serialbox)' is supported. It reads + from ppser serialized test data + Returns: IconGrid parsed from a given input type. + """ + if ser_type == SerializationType.SB: + return ( + sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) + .from_savepoint_grid() + .construct_icon_grid() + ) + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_initial_state( + gridfile_path: Path, rank=0 +) -> tuple[sb.IconSerialDataProvider, DiffusionDiagnosticState, PrognosticState]: + """ + Read prognostic and diagnostic state from serialized data. + + Args: + gridfile_path: path the serialized input data + + Returns: a tuple containing the data_provider, the initial diagnostic and prognostic state. + The data_provider is returned such that further timesteps of diagnostics and prognostics + can be read from within the dummy timeloop + + """ + data_provider = sb.IconSerialDataProvider( + "icon_pydycore", str(gridfile_path), False, mpi_rank=rank + ) + init_savepoint = data_provider.from_savepoint_diffusion_init( + linit=True, date=SIMULATION_START_DATE + ) + prognostic_state = init_savepoint.construct_prognostics() + diagnostic_state = init_savepoint.construct_diagnostics_for_diffusion() + return data_provider, diagnostic_state, prognostic_state + + +def read_geometry_fields( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> tuple[EdgeParams, CellParams, VerticalModelParams]: + """ + Read fields containing grid properties. + + Args: + path: path to the serialized input data + ser_type: (optional) defaults to SB=serialbox, type of input data to be read + + Returns: a tuple containing fields describing edges, cells, vertical properties of the model + the data is originally obtained from the grid file (horizontal fields) or some special input files. + """ + if ser_type == SerializationType.SB: + sp = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank + ).from_savepoint_grid() + edge_geometry = sp.construct_edge_geometry() + cell_geometry = sp.construct_cell_geometry() + vertical_geometry = VerticalModelParams(vct_a=sp.vct_a(), rayleigh_damping_height=12500) + return edge_geometry, cell_geometry, vertical_geometry + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_decomp_info( + path: Path, + procs_props: ProcessProperties, + ser_type=SerializationType.SB, +) -> DecompositionInfo: + if ser_type == SerializationType.SB: + sp = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), True, procs_props.rank + ) + return sp.from_savepoint_grid().construct_decomposition_info() + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_static_fields( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> tuple[DiffusionMetricState, DiffusionInterpolationState]: + """ + Read fields for metric and interpolation state. + + Args: + path: path to the serialized input data + rank: mpi rank, defaults to 0 for serial run + ser_type: (optional) defaults to SB=serialbox, type of input data to be read + + Returns: + a tuple containing the metric_state and interpolation state, + the fields are precalculated in the icon setup. + + """ + if ser_type == SerializationType.SB: + dataprovider = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank + ) + interpolation_state = ( + dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() + ) + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + return metric_state, interpolation_state + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def configure_logging(run_path: str, start_time, processor_procs: ProcessProperties = None) -> None: + """ + Configure logging. + + Log output is sent to console and to a file. + + Args: + run_path: path to the output folder where the logfile should be stored + start_time: start time of the model run + + """ + run_dir = Path(run_path).absolute() if run_path else Path(__file__).absolute().parent + run_dir.mkdir(exist_ok=True) + logfile = run_dir.joinpath(f"dummy_dycore_driver_{datetime.isoformat(start_time)}.log") + logfile.touch(exist_ok=True) + logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s %(filename)-20s (%(lineno)-4d) : %(funcName)-20s: %(levelname)-8s %(message)s", + filemode="w", + filename=logfile, + ) + console_handler = logging.StreamHandler() + console_handler.addFilter(ParallelLogger(processor_procs)) + + log_format = "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" + formatter = logging.Formatter(fmt=log_format, style="{", defaults={"rank": None}) + console_handler.setFormatter(formatter) + console_handler.setLevel(logging.DEBUG) + logging.getLogger("").addHandler(console_handler) diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py new file mode 100644 index 0000000000..d9ad439dc9 --- /dev/null +++ b/model/driver/tests/conftest.py @@ -0,0 +1,24 @@ +# 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 icon4py.model.common.test_utils.fixtures import ( # noqa: F401 + datapath, + download_ser_data, + processor_props, + ranked_data_path, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/driver/tests/test_io_utils.py b/model/driver/tests/test_io_utils.py new file mode 100644 index 0000000000..de51e2af1a --- /dev/null +++ b/model/driver/tests/test_io_utils.py @@ -0,0 +1,108 @@ +# 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 + + +import pytest + +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.driver.io_utils import ( + SerializationType, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) + + +@pytest.mark.datatest +@pytest.mark.parametrize("read_fun", (read_geometry_fields, read_static_fields, read_icon_grid)) +def test_read_geometry_fields_not_implemented_type(read_fun, datapath): + with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): + read_fun(path=datapath, ser_type=SerializationType.NC) + + +def assert_grid_size_and_connectivities(grid): + assert grid.num_edges() == 31558 + assert grid.num_cells() == 20896 + assert grid.num_vertices() == 10663 + assert grid.get_e2v_connectivity() + assert grid.get_v2e_connectivity() + assert grid.get_c2e_connectivity() + assert grid.get_e2c_connectivity() + assert grid.get_e2c2v_connectivity() + assert grid.get_c2e2c_connectivity() + assert grid.get_e2ecv_connectivity() + + +@pytest.mark.datatest +def test_read_icon_grid_for_type_sb(datapath): + grid = read_icon_grid(datapath, ser_type=SerializationType.SB) + assert_grid_size_and_connectivities(grid) + + +@pytest.mark.datatest +def test_read_static_fields_for_type_sb(datapath): + metric_state, interpolation_state = read_static_fields(datapath, ser_type=SerializationType.SB) + assert_metric_state_fields(metric_state) + assert_interpolation_state_fields(interpolation_state) + + +@pytest.mark.datatest +def test_read_geometry_fields_for_type_sb(datapath): + edge_geometry, cell_geometry, vertical_geometry = read_geometry_fields( + datapath, ser_type=SerializationType.SB + ) + assert_edge_geometry_fields(edge_geometry) + assert_cell_geometry_fields(cell_geometry) + assert_vertical_params(vertical_geometry) + + +def assert_vertical_params(vertical_geometry: VerticalModelParams): + assert vertical_geometry.physical_heights + assert vertical_geometry.index_of_damping_layer > 0 + assert vertical_geometry.rayleigh_damping_height > 0 + + +def assert_cell_geometry_fields(cell_geometry: CellParams): + assert cell_geometry.area + + +def assert_edge_geometry_fields(edge_geometry: EdgeParams): + assert edge_geometry.edge_areas + assert edge_geometry.primal_normal_vert + assert edge_geometry.inverse_primal_edge_lengths + assert edge_geometry.tangent_orientation + assert edge_geometry.inverse_dual_edge_lengths + assert edge_geometry.dual_normal_vert + + +def assert_metric_state_fields(metric_state): + assert metric_state.wgtfac_c + assert metric_state.zd_intcoef + assert metric_state.zd_diffcoef + assert metric_state.theta_ref_mc + assert metric_state.mask_hdiff + assert metric_state.zd_vertoffset + + +def assert_interpolation_state_fields(interpolation_state): + assert interpolation_state.geofac_n2s + assert interpolation_state.e_bln_c_s + assert interpolation_state.nudgecoeff_e + assert interpolation_state.geofac_n2s_nbh + assert interpolation_state.geofac_div + assert interpolation_state.geofac_grg_y + assert interpolation_state.geofac_grg_x + assert interpolation_state.rbf_coeff_2 + assert interpolation_state.rbf_coeff_1 + assert interpolation_state.geofac_n2s_c diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index 10bde64b68..d4c4c5727f 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -1,4 +1,6 @@ -r ../base-requirements-dev.txt -e ./atmosphere/dycore +-e ./atmosphere/diffusion -e ./atmosphere/advection -e ./common +-e ./driver diff --git a/model/requirements.txt b/model/requirements.txt index 4406060faf..3330309e2e 100644 --- a/model/requirements.txt +++ b/model/requirements.txt @@ -1,4 +1,6 @@ -r ../base-requirements.txt ./atmosphere/dycore +./atmosphere/diffusion ./atmosphere/advection ./common +./driver diff --git a/model/tox.ini b/model/tox.ini index e14438ea04..09c7bbca4a 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/advection/src common/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src atmosphere/advection/src common/src driver/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html @@ -31,3 +31,4 @@ setenv = skip_install = true commands = commands_post = + diff --git a/requirements-dev.txt b/requirements-dev.txt index 809c77a57b..61f75ec630 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -4,6 +4,8 @@ -e ./model/atmosphere/dycore -e ./model/atmosphere/advection -e ./model/common +-e ./model/atmosphere/diffusion +-e ./model/driver # icon4pytools -e ./tools diff --git a/requirements.txt b/requirements.txt index ca7931901c..656d6c7d29 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,10 @@ # icon4py model ./model/atmosphere/dycore +./model/atmosphere/diffusion ./model/atmosphere/advection ./model/common +./model/driver # icon4pytools ./tools diff --git a/tools/README.md b/tools/README.md index fdb881588a..41f040e0dc 100644 --- a/tools/README.md +++ b/tools/README.md @@ -232,3 +232,39 @@ OUTPUT_FILEPATH A path to the output Fortran source file to be generated. ``` **Note:** The output of f2ser still has to be preprocessed using `pp_ser.py`, which then yields a compilable unit. The serialised files will have `f2ser` as their prefix in the default folder location of the experiment. + +## `py2f` + +Python utility for generating a C library and Fortran interface to call Python icon4py modules. The library [embeds python via CFFI ](https://cffi.readthedocs.io/en/latest/embedding.) + +This is **highly experimental** and has not been tested from within Fortran code! + +### usage + +`py2fgen`: Generates a C header file and a Fortran interface and compiles python functions into a C library embedding python. The functions need to be decorated with `CffiMethod.register` and have a signature with scalar arguments or `GT4Py` fields, for example: + +``` +@CffiMethod.register +def foo(i:int, param:float, field1: Field[[VertexDim, KDim], float], field2: Field[CellDim, KDim], float]) +``` + +see `src/icon4pytools/py2f/wrappers` for examples. + +```bash +py2fgen icon4pytools.py2f.wrappers.diffusion_wrapper py2f_build +``` + +where the first argument is the python module to parse and the second a build directory. The call above will generate the following in `py2f_build`: + +```bash + ls py2f_build/ +total 204K +drwxrwxr-x 2 magdalena magdalena 4.0K Aug 24 17:01 . +drwxrwxr-x 9 magdalena magdalena 4.0K Aug 24 17:01 .. +-rw-rw-r-- 1 magdalena magdalena 74K Aug 24 16:58 diffusion_wrapper.c +-rw-rw-r-- 1 magdalena magdalena 3.5K Aug 24 17:01 diffusion_wrapper.f90 +-rw-rw-r-- 1 magdalena magdalena 955 Aug 24 17:01 diffusion_wrapper.h +-rw-rw-r-- 1 magdalena magdalena 58K Aug 24 17:01 diffusion_wrapper.o +-rwxrwxr-x 1 magdalena magdalena 49K Aug 24 17:01 libdiffusion_wrapper.so + +``` diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 9faa971e41..06ee7aa9be 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -19,12 +19,14 @@ classifiers = [ 'Topic :: Scientific/Engineering :: Physics' ] dependencies = [ + 'icon4py-atmosphere-diffusion', 'icon4py-atmosphere-dycore', 'icon4py-atmosphere-advection', 'gt4py>=1.0.1', - 'icon4py_common', + 'icon4py-common', 'tabulate>=0.8.9', - 'fprettify>=0.3.7' + 'fprettify>=0.3.7', + 'cffi>=1.5' ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] @@ -37,6 +39,7 @@ requires-python = '>=3.10' f2ser = 'icon4pytools.f2ser.cli:main' icon4pygen = 'icon4pytools.icon4pygen.cli:main' icon_liskov = 'icon4pytools.liskov.cli:main' +py2fgen = 'icon4pytools.py2f.py2fgen:main' [project.urls] repository = 'https://github.com/C2SM/icon4py' @@ -109,7 +112,12 @@ use_parentheses = true [tool.mypy] disallow_incomplete_defs = true disallow_untyped_defs = true -exclude = ['^tests/f2ser/*.py', '^tests/icon4pygen/*.py', '^tests/liskov/*.py'] +exclude = [ + '^tests/f2ser/*.py', + '^tests/icon4pygen/*.py', + '^tests/liskov/*.py', + '^tests/py2f/*.py' +] ignore_missing_imports = false implicit_reexport = true install_types = true diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index 4dc72f7d88..16afdad85e 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,5 +1,6 @@ -r ../base-requirements-dev.txt -e ../model/atmosphere/dycore +-e ../model/atmosphere/diffusion -e ../model/atmosphere/advection -e ../model/common -e . diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 53b815173f..6edcc302a0 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -32,7 +32,11 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atmosphere.dycore"] + _STENCIL_PACKAGES = [ + "atmosphere.dycore", + "atmosphere.diffusion.stencils", + "common.interpolation.stencils", + ] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed diff --git a/tools/src/icon4pytools/py2f/__init__.py b/tools/src/icon4pytools/py2f/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/tools/src/icon4pytools/py2f/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py new file mode 100644 index 0000000000..c251831322 --- /dev/null +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -0,0 +1,177 @@ +# 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 +import functools +import inspect +from collections import OrderedDict +from importlib.resources import files +from types import MappingProxyType +from typing import Any + +import cffi +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4pytools.py2f.typing_utils import parse_annotation + + +FFI_DEF_EXTERN_DECORATOR = "@ffi.def_extern()" + +CFFI_GEN_DECORATOR = "@CffiMethod.register" + + +class CffiMethod: + _registry = {} + + @classmethod + def register(cls, func): + cls._registry.setdefault(func.__module__, []).append(func.__name__) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + + @classmethod + def get(cls, name: str): + return cls._registry[name] + + +def generate_and_compile_cffi_plugin( + plugin_name: str, c_header: str, module_name: str, build_path="." +): + """ + Create C shared library. + + Create a linkable C library and F90 interface for the functions in the python module + {module_name} that are decorated with '@CffiMethod.register'. + + Args: + plugin_name: name of the plugin, a linkable C library with the name + 'lib{plugin_name}.so' will be created in the {build_path} folder' + c_header: C type header signature for the python functions. + module_name: python module name that contains python functions corresponding + to the signature in the '{c_header}' string, these functions must be decorated + with @CffiMethod.register and the file must contain the import + build_path: *optional* path to build directory + + """ + module_split = module_name.split(".") + python_src_file = "/".join(module_split[1:]) + ".py" + python_package = module_split[0] + + c_header_file = plugin_name + ".h" + with open("/".join([build_path, c_header_file]), "w") as f: + f.write(c_header) + + builder = cffi.FFI() + + builder.embedding_api(c_header) + builder.set_source(plugin_name, f'#include "{c_header_file}"') + + module = files(python_package).joinpath(python_src_file).read_text() + + module = f"from {plugin_name} import ffi\n{module}".replace( + CFFI_GEN_DECORATOR, FFI_DEF_EXTERN_DECORATOR + ) + + builder.embedding_init_code(module) + builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) + + +class UnknownDimensionException(Exception): + """Raised if a Dimension is unknown to the interface generation.""" + + pass + + +def to_fields(dim_sizes: dict[Dimension, int]): + """ + Pack/Unpack Fortran 2d arrays to numpy arrays with using CFFI frombuffer. + + Args: + dim_sizes: dictionary containing the sizes of the dimension. + + #TODO (magdalena) handle dimension sizes in a better way? + """ + ffi = cffi.FFI() + dim_sizes = dim_sizes + + def _dim_sizes(dims: list[Dimension]) -> tuple[int, int]: + """Extract the size of dimension from a dictionary.""" + v_size = None + h_size = None + for d in dims: + if d not in dim_sizes.keys(): + raise UnknownDimensionException( + f"size of dimension '{d}' not defined in '{dim_sizes.keys()}'" + ) + if d.kind == DimensionKind.VERTICAL or d.kind == DimensionKind.LOCAL: + v_size = dim_sizes[d] + elif d.kind == DimensionKind.HORIZONTAL: + h_size = dim_sizes[d] + return h_size, v_size + + def _unpack(ptr, size_h, size_v, dtype) -> np.ndarray: + """ + Unpack a 2d c/fortran field into a numpy array. + + :param dtype: expected type of the fields + :param ptr: c_pointer to the field + :param size_h: length of horizontal dimension + :param size_v: length of vertical dimension + :return: a numpy array with shape=(size_h, size_v) + and dtype = ctype of the pointer + """ + shape = (size_h, size_v) + length = np.prod(shape) + c_type = ffi.getctype(ffi.typeof(ptr).item) + # TODO (magdalena) fix dtype handling use SCALARTYPE? + mem_size = ffi.sizeof(c_type) + mem_size = np.dtype(c_type).itemsize + ar = np.frombuffer( + ffi.buffer(ptr, length * mem_size), + dtype=np.dtype(c_type), + count=-1, + offset=0, + ).reshape(shape) + return ar + + def _to_fields_decorator(func): + @functools.wraps(func) + def _wrapper( + *args, **kwargs + ): # these are the args of the decorated function ie to_fields(func(*args, **kwargs)) + signature = inspect.signature(func) + parameters: MappingProxyType[str, inspect.Parameter] = signature.parameters + arguments: OrderedDict[str, Any] = signature.bind(*args, **kwargs).arguments + f_args = [] + for name, argument in arguments.items(): + ar = _transform_arg(argument, name, parameters) + f_args.append(ar) + return func(*f_args, **kwargs) + + def _transform_arg(argument, name, parameters): + dims, dtype = parse_annotation(parameters[name].annotation) + if dims: + (size_h, size_v) = _dim_sizes(dims) + ar = _unpack(argument, size_h, size_v, dtype) + ar = np_as_located_field(*dims)(ar) + else: + ar = argument + return ar + + return _wrapper + + return _to_fields_decorator diff --git a/tools/src/icon4pytools/py2f/codegen.py b/tools/src/icon4pytools/py2f/codegen.py new file mode 100644 index 0000000000..502b204922 --- /dev/null +++ b/tools/src/icon4pytools/py2f/codegen.py @@ -0,0 +1,145 @@ +# 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 typing import Sequence + +from gt4py.eve import Node, codegen +from gt4py.eve.codegen import JinjaTemplate as as_jinja +from gt4py.eve.codegen import TemplatedGenerator +from gt4py.next.type_system.type_specifications import ScalarKind + +from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( + BUILTIN_TO_CPP_TYPE, + BUILTIN_TO_ISO_C_TYPE, +) +from icon4pytools.icon4pygen.bindings.utils import write_string + + +class DimensionType(Node): + name: str + length: int + + +class FuncParameter(Node): + name: str + d_type: ScalarKind + dimensions: Sequence[DimensionType] + + +class Func(Node): + name: str + args: Sequence[FuncParameter] + + +class CffiPlugin(Node): + name: str + functions: Sequence[Func] + + +def to_c_type(scalar_type: ScalarKind) -> str: + return BUILTIN_TO_CPP_TYPE[scalar_type] + + +def to_f_type(scalar_type: ScalarKind) -> str: + return BUILTIN_TO_ISO_C_TYPE[scalar_type] + + +def as_f90_value(param: FuncParameter) -> str: + """ + If param is a scalar type (dimension=0) then return the F90 'value' keyword. + + Used for F90 generation only. + """ + return "value, " if len(param.dimensions) == 0 else "" + + +def as_field(param: FuncParameter, language: str) -> str: + size = len(param.dimensions) + if size == 0: + return "" + if "C" == language: + return "*" + else: + dims = ",".join(map(lambda x: ":", range(size))) + return f"({dims})" + + +class CHeaderGenerator(TemplatedGenerator): + CffiPlugin = as_jinja("""{% for func in functions: %}{{func}}\n{% endfor %}""") + + Func = as_jinja("""extern void {{name}}({{", ".join(args)}});""") + + def visit_FuncParameter(self, param: FuncParameter): + return self.generic_visit( + param, + rendered_type=to_c_type(param.d_type), + dim=as_field(param, "C"), + ) + + FuncParameter = as_jinja("""{{rendered_type}}{{dim}} {{name}}""") + + +class F90InterfaceGenerator(TemplatedGenerator): + CffiPlugin = as_jinja( + """ + module {{name}} + use, intrinsic:: iso_c_binding + implicit none + + public + interface + {% for func in functions: %}\ + {{func}}\ + {% endfor %}\ + end interface + end module + """ + ) + + def visit_Func(self, func: Func): + arg_names = ", ".join(map(lambda x: x.name, func.args)) + return self.generic_visit(func, param_names=arg_names) + + Func = as_jinja( + """subroutine {{name}}({{param_names}}) bind(c, name='{{name}}') + use iso_c_binding + {% for arg in args: %}\ + {{arg}}\ + {% endfor %}\ + end subroutine {{name}} + """ + ) + + def visit_FuncParameter(self, param: FuncParameter, param_names=""): + # kw-arg param_names needs to be present because it is present up the tree + return self.generic_visit( + param, + value=as_f90_value(param), + rendered_type=to_f_type(param.d_type), + dim=as_field(param, "F"), + ) + + FuncParameter = as_jinja( + """{{rendered_type}}, {{value}} intent(inout):: {{name}}{{dim}} + """ + ) + + +def generate_c_header(plugin: CffiPlugin) -> str: + generated_code = CHeaderGenerator.apply(plugin) + return codegen.format_source("cpp", generated_code, style="LLVM") + + +def generate_and_write_f90_interface(build_path: str, plugin: CffiPlugin): + generated_code = F90InterfaceGenerator.apply(plugin) + write_string(generated_code, build_path, f"{plugin.name}.f90") diff --git a/tools/src/icon4pytools/py2f/parsing.py b/tools/src/icon4pytools/py2f/parsing.py new file mode 100644 index 0000000000..847285615f --- /dev/null +++ b/tools/src/icon4pytools/py2f/parsing.py @@ -0,0 +1,43 @@ +# 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 + +import importlib +from inspect import signature, unwrap + +from icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter +from icon4pytools.py2f.typing_utils import parse_annotation + + +def parse_functions_from_module(module_name: str) -> CffiPlugin: + module = importlib.import_module(module_name) + func_names = CffiMethod.get(module_name) + funcs = [_parse_function(module, fn) for fn in func_names] + plugin_name = module_name.split(".")[-1] + return CffiPlugin(name=plugin_name, functions=funcs) + + +def _parse_function(module, s): + func = unwrap(getattr(module, s)) + params = [ + _parse_params(signature(func, follow_wrapped=False).parameters, p) + for p in (signature(func).parameters) + ] + return Func(name=s, args=params) + + +def _parse_params(params, s): + annotation = params[s].annotation + dims, dtype = parse_annotation(annotation) + dim_types = [DimensionType(name=d.value, length=10) for d in dims] + return FuncParameter(name=s, d_type=dtype, dimensions=dim_types) diff --git a/tools/src/icon4pytools/py2f/py2fgen.py b/tools/src/icon4pytools/py2f/py2fgen.py new file mode 100644 index 0000000000..7bdb799d1b --- /dev/null +++ b/tools/src/icon4pytools/py2f/py2fgen.py @@ -0,0 +1,51 @@ +# 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 + +import pathlib + +import click + +from icon4pytools.py2f.cffi_utils import generate_and_compile_cffi_plugin +from icon4pytools.py2f.codegen import generate_and_write_f90_interface, generate_c_header +from icon4pytools.py2f.parsing import parse_functions_from_module + + +@click.command( + "py2fgen", +) +@click.argument("module", type=str) +@click.argument( + "build_path", + type=click.Path(dir_okay=True, resolve_path=True, path_type=pathlib.Path), + default=".", +) +def main(module: str, build_path: pathlib.Path) -> None: + """ + Generate C and F90 wrappers and C library for embedding the python MODULE in C and Fortran. + + Args: + - module: name of the python module containing the methods to be embedded. Those + methods have to be decorated with CffiMethod.register + + - build_path: directory where the generated code and compiled libraries are to be found. + """ + module_name = module + build_path.mkdir(exist_ok=True, parents=True) + plugin = parse_functions_from_module(module_name) + c_header = generate_c_header(plugin) + generate_and_compile_cffi_plugin(plugin.name, c_header, module_name, str(build_path)) + generate_and_write_f90_interface(build_path, plugin) + + +if __name__ == "__main__": + main() diff --git a/tools/src/icon4pytools/py2f/typing_utils.py b/tools/src/icon4pytools/py2f/typing_utils.py new file mode 100644 index 0000000000..72ba5fa286 --- /dev/null +++ b/tools/src/icon4pytools/py2f/typing_utils.py @@ -0,0 +1,27 @@ +# 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.next.common import Dimension +from gt4py.next.type_system.type_specifications import ScalarType +from gt4py.next.type_system.type_translation import from_type_hint + + +def parse_annotation(annotation) -> tuple[list[Dimension], ScalarType]: + type_spec = from_type_hint(annotation) + if isinstance(type_spec, ScalarType): + dtype = type_spec.kind + dims = [] + else: + dtype = type_spec.dtype.kind + dims = type_spec.dims + return dims, dtype diff --git a/tools/src/icon4pytools/py2f/wrappers/__init__.py b/tools/src/icon4pytools/py2f/wrappers/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/tools/src/icon4pytools/py2f/wrappers/__init__.py @@ -0,0 +1,12 @@ +# 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 diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py new file mode 100644 index 0000000000..89c2dbb59e --- /dev/null +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -0,0 +1,265 @@ +# 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 +# flake8: noqa +""" +Wrapper module for diffusion granule. + +Module contains a diffusion_init and diffusion_run function that follow the architecture of +Fortran granule interfaces: +- all arguments needed from external sources are passed. +- passing of scalar types or fields of simple types + + +""" +import numpy as np +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionConfig, + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + DiffusionParams, + DiffusionType, + PrognosticState, +) +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CECDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid +from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams + +from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields + + +# TODO (magdalena) Revise interface architecture with Fortran granules: +# The module variable to match the Fortran interface: where only fields are passed. +# We should rather instantiate the object init and return it. +diffusion: Diffusion() + +nproma = 50000 +field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} + + +@to_fields(dim_sizes=field_sizes) +@CffiMethod.register +def diffusion_init( + nproma: int, + nlev: int, + n_shift_total: int, + nrdmax: float, + n_dyn_substeps: int, + nudge_max_coeff: float, + denom_diffu_v: float, + hdiff_smag_fac: float, + hdiff_order: int, + hdiff_efdt_ratio: float, + itype_sher: int, + itype_vn_diffu: int, + itype_t_diffu: int, + lhdiff_rcf: bool, + lhdiff_w: bool, + lhdiff_temp: bool, + l_limited_area: bool, + l_zdiffu_t: bool, + lsmag_3d: bool, + vct_a: Field[[KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + nudgecoeff_e: Field[[EdgeDim], float], + zd_intcoef: Field[ + [CellDim, C2E2CDim, KDim], float + ], # special DSL field: zd_intcoef_dsl in mo_vertical_grid.f90 + zd_diffcoef: Field[ + [CellDim, KDim], float + ], # special DSL field mask instead of list: zd_diffcoef_dsl in mo_vertical_grid.f90 + wgtfac_c: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + edges_vertex_idx: Field[[EdgeDim, E2VDim], int32], + edges_cell_idx: Field[[EdgeDim, E2CDim], int32], + edges_tangent_orientation: Field[[EdgeDim], float], + edges_primal_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_primal_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_inv_vert_vert_lengths: Field[[EdgeDim], float], + edges_inv_primal_edge_length: Field[[EdgeDim], float], + edges_inv_dual_edge_length: Field[[EdgeDim], float], + edges_area_edge: Field[[EdgeDim], float], + cells_neighbor_idx: Field[[CellDim, C2E2CDim], int32], + cells_edge_idx: Field[[CellDim, C2EDim], int32], + cells_area: Field[[CellDim], float], + # dsl specific additional args + mask_hdiff: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[ + [CECDim], int32 + ], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 + rbf_coeff_1: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + rbf_coeff_2: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + verts_edge_idx: Field[[VertexDim, V2EDim], int32], # -> mo_intp_rbf_rbf_vec_interpol_vertex +): + """ + Instantiate and Initialize the diffusion object. + + should only accept simple fields as arguments for compatibility with the standalone + Fortran ICON Diffusion component (aka Diffusion granule) + + """ + if diffusion.initialized: + raise DuplicateInitializationException("Diffusion has already been already initialized") + + horizontal_size = HorizontalGridSize(num_cells=nproma, num_vertices=nproma, num_edges=nproma) + vertical_size = VerticalGridSize(num_lev=nlev) + mesh_config = GridConfig( + horizontal_config=horizontal_size, + vertical_config=vertical_size, + limited_area=l_limited_area, + n_shift_total=n_shift_total, + ) + # we need the start, end indices in order for the grid to be functional those are not passed + # to init in the Fortran diffusion granule since they are hidden away in the get_indices_[c,e,v] + # diffusion_run does take p_patch in order to pass it on to other subroutines (interpolation, get_indices... + + c2e2c0 = np.column_stack( + ( + (np.asarray(cells_neighbor_idx)), + (np.asarray(range(np.asarray(cells_neighbor_idx).shape[0]))), + ) + ) + e2c2v = np.asarray(edges_vertex_idx) + e2v = e2c2v[:, 0:2] + connectivities = { + E2CDim: np.asarray(edges_cell_idx), + E2C2VDim: e2c2v, + E2VDim: e2v, + C2EDim: np.asarray(cells_edge_idx), + C2E2CDim: (np.asarray(cells_neighbor_idx)), + C2E2CODim: c2e2c0, + V2EDim: np.asarray(verts_edge_idx), + } + # TODO (Magdalena) we need start_index, end_index: those are not passed in the fortran granules, + # because they are used through get_indices only + grid = IconGrid().with_config(mesh_config).with_connectivities(connectivities) + edge_params = EdgeParams( + tangent_orientation=edges_tangent_orientation, + inverse_primal_edge_lengths=edges_inv_primal_edge_length, + dual_normal_vert_x=edges_dual_normal_vert_1, + dual_normal_vert_y=edges_dual_normal_vert_2, + inverse_dual_edge_lengths=edges_inv_dual_edge_length, + inverse_vertex_vertex_lengths=edges_inv_vert_vert_lengths, + primal_normal_vert_x=edges_primal_normal_vert_1, + primal_normal_vert_y=edges_primal_normal_vert_2, + edge_areas=edges_area_edge, + ) + cell_params = CellParams(cells_area) + config: DiffusionConfig = DiffusionConfig( + diffusion_type=DiffusionType(hdiff_order), + hdiff_w=lhdiff_w, + hdiff_temp=lhdiff_temp, + type_vn_diffu=itype_vn_diffu, + smag_3d=lsmag_3d, + type_t_diffu=itype_t_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + hdiff_w_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=hdiff_smag_fac, + n_substeps=n_dyn_substeps, + zdiffu_t=l_zdiffu_t, + hdiff_rcf=lhdiff_rcf, + velocity_boundary_diffusion_denom=denom_diffu_v, + max_nudging_coeff=nudge_max_coeff, + ) + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) + + derived_diffusion_params = DiffusionParams(config) + metric_state = DiffusionMetricState( + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + mask_hdiff=mask_hdiff, + zd_vertoffset=zd_vertoffset, + zd_diffcoef=zd_diffcoef, + zd_intcoef=zd_intcoef, + ) + interpolation_state = DiffusionInterpolationState( + e_bln_c_s, + rbf_coeff_1, + rbf_coeff_2, + geofac_div, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + nudgecoeff_e, + ) + + diffusion.init( + grid=grid, + config=config, + params=derived_diffusion_params, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + +@CffiMethod.register +def diffusion_run( + dtime: float, + linit: bool, + vn: Field[[EdgeDim, KDim], float], + w: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], +): + diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + prognostic_state = PrognosticState( + w=w, + vn=vn, + exner_pressure=exner, + theta_v=theta_v, + ) + if linit: + diffusion.initial_run( + diagnostic_state, + prognostic_state, + dtime, + ) + else: + diffusion.run(diagnostic_state, prognostic_state, dtime) + + +class DuplicateInitializationException(Exception): + """Raised if the component is already initilalized.""" + + pass diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 4ca1a1d1f6..72c35d898b 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -15,7 +15,9 @@ import pkgutil import re +import icon4py.model.atmosphere.diffusion.stencils as diffusion import icon4py.model.atmosphere.dycore as dycore +import icon4py.model.common.interpolation.stencils as intp import pytest from click.testing import CliRunner @@ -25,6 +27,9 @@ DYCORE_PKG = "atmosphere.dycore" +INTERPOLATION_PKG = "common.interpolation.stencils" +DIFFUSION_PKG = "atmosphere.diffusion.stencils" + LEVELS_PER_THREAD = "1" BLOCK_SIZE = "128" OUTPATH = "." @@ -36,9 +41,21 @@ def cli(): def dycore_fencils() -> list[tuple[str, str]]: - pkgpath = os.path.dirname(dycore.__file__) + return _fencils(dycore.__file__, DYCORE_PKG) + + +def interpolation_fencils() -> list[tuple[str, str]]: + return _fencils(intp.__file__, INTERPOLATION_PKG) + + +def diffusion_fencils() -> list[tuple[str, str]]: + return _fencils(diffusion.__file__, DIFFUSION_PKG) + + +def _fencils(module_name, package_name) -> list[tuple[str, str]]: + pkgpath = os.path.dirname(module_name) stencils = [name for _, name, _ in pkgutil.iter_modules([pkgpath])] - fencils = [(DYCORE_PKG, stencil) for stencil in stencils] + fencils = [(package_name, stencil) for stencil in stencils] return fencils @@ -111,7 +128,10 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") -@pytest.mark.parametrize(("stencil_module", "stencil_name"), dycore_fencils()) +@pytest.mark.parametrize( + ("stencil_module", "stencil_name"), + dycore_fencils() + interpolation_fencils() + diffusion_fencils(), +) def test_codegen_dycore(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) with cli.isolated_filesystem(): diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 739c2ba217..55f2e5f3b2 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -40,7 +40,7 @@ def test_stencil_collector_invalid_module(): def test_stencil_collector_invalid_member(): - from icon4py.model.atmosphere.dycore import apply_nabla2_to_w + from icon4py.model.atmosphere.diffusion.stencils import apply_nabla2_to_w module_path = Path(apply_nabla2_to_w.__file__) parents = module_path.parents[0] diff --git a/tools/tests/py2f/test_cffi_utils.py b/tools/tests/py2f/test_cffi_utils.py new file mode 100644 index 0000000000..d8fdf61968 --- /dev/null +++ b/tools/tests/py2f/test_cffi_utils.py @@ -0,0 +1,80 @@ +# 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 + +import cffi +import numpy as np +import pytest +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import float32, float64, int32, int64 +from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim, VertexDim + +from icon4pytools.py2f.cffi_utils import UnknownDimensionException, to_fields + + +n_vertices = 9 +n_edges = 27 +levels = 12 +e2c_sparse_size = 2 + + +def random(*sizes): + return np.random.default_rng().uniform(size=sizes) + + +@pytest.mark.parametrize("pointer_type", ["float*", "double*"]) +def test_unpack_from_buffer_to_field(pointer_type: str): + @to_fields(dim_sizes={VertexDim: n_vertices, KDim: levels}) + def identity( + f1: Field[[VertexDim, KDim], float], factor: float + ) -> tuple[float, Field[[VertexDim, KDim], float]]: + return factor, f1 + + ffi = cffi.FFI() + input_array = random(n_vertices, levels) + input_factor = 0.5 + res_factor, result_field = identity(ffi.from_buffer(pointer_type, input_array), input_factor) + assert res_factor == input_factor + assert np.allclose(np.asarray(result_field), input_array) + + +def test_unpack_only_scalar_args(): + @to_fields(dim_sizes={}) + def multiply(f1: float, f2: float32, f3: float64, i3: int64, i2: int32, i1: int): + return f1 * f2 * f3, i1 * i2 * i3 + + f_res, i_res = multiply(0.5, 2.0, 3.0, 1, 2, 3) + assert f_res == 3.0 + assert i_res == 6 + + +@pytest.mark.parametrize("field_type", [int32, int, int64]) +def test_unpack_local_field(field_type): + ffi = cffi.FFI() + + @to_fields(dim_sizes={EdgeDim: n_edges, E2CDim: e2c_sparse_size}) + def local_field(f1: Field[[EdgeDim, E2CDim], field_type]): + return f1 + + input_field = np.arange(n_edges * e2c_sparse_size).reshape((n_edges, e2c_sparse_size)) + res_field = local_field(ffi.from_buffer("int*", input_field)) + assert np.all(res_field == input_field) + + +def test_unknown_dimension_raises_exception(): + @to_fields(dim_sizes={}) + def do_nothing(f1: Field[[VertexDim], float]): + pass + + input_array = random() + with pytest.raises(UnknownDimensionException, match=r"size of dimension "): + do_nothing(input_array) diff --git a/tools/tests/py2f/test_code_generation.py b/tools/tests/py2f/test_code_generation.py new file mode 100644 index 0000000000..74a9326de2 --- /dev/null +++ b/tools/tests/py2f/test_code_generation.py @@ -0,0 +1,140 @@ +# 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 +import string + +import pytest +from gt4py.next.type_system.type_specifications import ScalarKind + +from icon4pytools.py2f.codegen import ( + CffiPlugin, + CHeaderGenerator, + DimensionType, + F90InterfaceGenerator, + Func, + FuncParameter, + as_f90_value, + as_field, +) + + +field_2d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13), DimensionType(name="J", length=13)], +) +field_1d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13)], +) + +simple_type = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) + + +@pytest.mark.parametrize( + ("param", "expected"), ((simple_type, "value, "), (field_2d, ""), (field_1d, "")) +) +def test_as_target(param, expected): + assert expected == as_f90_value(param) + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:,:)"))) +def test_field_extension_2d(lang, expected): + assert as_field(field_2d, lang) == expected + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:)"))) +def test_field_extension_1d(lang, expected): + assert as_field(field_1d, lang) == expected + + +@pytest.mark.parametrize("lang", ("C", "F")) +def test_is_field_simple_type(lang): + assert as_field(simple_type, lang) == "" + + +foo = Func( + name="foo", + args=[ + FuncParameter(name="one", d_type=ScalarKind.INT32, dimensions=[]), + FuncParameter(name="two", d_type=ScalarKind.FLOAT64, dimensions=[]), + ], +) + +bar = Func( + name="bar", + args=[ + FuncParameter( + name="one", + d_type=ScalarKind.FLOAT32, + dimensions=[ + DimensionType(name="KDim", length=10), + DimensionType(name="VDim", length=50000), + ], + ), + FuncParameter(name="two", d_type=ScalarKind.INT32, dimensions=[]), + ], +) + + +def test_cheader_generation_for_single_function(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void foo(int one, double two);\n" + + +def test_cheader_for_pointer_args(): + functions = [bar] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void bar(float* one, int two);\n" + + +def compare_ignore_whitespace(s1: str, s2: str): + no_whitespace = {ord(c): None for c in string.whitespace} + return s1.translate(no_whitespace) == s2.translate(no_whitespace) + + +def test_c_header_with_several_functions(): + functions = [bar, foo] + plugin = CffiPlugin(name="libtest", functions=functions) + header = CHeaderGenerator.apply(plugin) + assert ( + header + == """extern void bar(float* one, int two);\nextern void foo(int one, double two);\n""" + ) + + +def test_fortran_interface(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + interface = F90InterfaceGenerator.apply(plugin) + expected = """ + module libtest + use, intrinsic:: iso_c_binding + implicit none + + public + interface + subroutine foo(one, two) bind(c, name='foo') + use iso_c_binding + integer(c_int), value, intent(inout):: one + real(c_double), value, intent(inout):: two + end subroutine foo + end interface + end module + """ + assert compare_ignore_whitespace(interface, expected) diff --git a/tools/tests/py2f/test_parsing_wrapper.py b/tools/tests/py2f/test_parsing_wrapper.py new file mode 100644 index 0000000000..ea19520384 --- /dev/null +++ b/tools/tests/py2f/test_parsing_wrapper.py @@ -0,0 +1,39 @@ +# 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 icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.parsing import parse_functions_from_module + + +def test_parse_functions(): + path = "icon4pytools.py2f.wrappers.diffusion_wrapper" + plugin = parse_functions_from_module(path) + + assert plugin.name == "diffusion_wrapper" + assert len(plugin.functions) == 2 + assert "diffusion_init" in map(lambda f: f.name, plugin.functions) + assert "diffusion_run" in map(lambda f: f.name, plugin.functions) + + +@CffiMethod.register +def do_foo(foo: str): + return foo + + +@CffiMethod.register +def do_bar(): + return "bar" + + +def test_register_with_cffi(): + assert "do_foo" in CffiMethod.get(__name__) + assert "do_bar" in CffiMethod.get(__name__) diff --git a/tools/tests/py2f/test_py2fgen.py b/tools/tests/py2f/test_py2fgen.py new file mode 100644 index 0000000000..8330e4a267 --- /dev/null +++ b/tools/tests/py2f/test_py2fgen.py @@ -0,0 +1,25 @@ +# 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 click.testing import CliRunner + +from icon4pytools.py2f.py2fgen import main + + +def test_py2fgen(): + cli = CliRunner() + module = "icon4pytools.py2f.wrappers.diffusion_wrapper" + build_path = "./build" + with cli.isolated_filesystem(): + result = cli.invoke(main, [module, build_path]) + assert result.exit_code == 0 diff --git a/tox.ini b/tox.ini index eec85c5901..d32caf0541 100644 --- a/tox.ini +++ b/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/advection/src model/common/src tools/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/atmosphere/advection/src model/driver/src model/common/src tools/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html From 08724557e2fe63b8cbd5dcf9d05c6d4194dbef44 Mon Sep 17 00:00:00 2001 From: juckerj <39263956+jonasjucker@users.noreply.github.com> Date: Wed, 13 Sep 2023 11:49:11 +0200 Subject: [PATCH 096/105] [CI] add option to use custom spack-c2sm branch for spack tests (#268) * add option to use custom spack-c2sm branch for spack tests * fix missing declarations * fix missing declarations --------- Co-authored-by: Jonas Jucker --- jenkins/spack-PR | 17 ++++++++++++++++- jenkins/spack-PR-icon | 5 ----- jenkins/spack-PR-stable | 17 ++++++++++++++++- 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/jenkins/spack-PR b/jenkins/spack-PR index 322dd30209..e9be0ac87c 100644 --- a/jenkins/spack-PR +++ b/jenkins/spack-PR @@ -1,3 +1,18 @@ +def repo_identifier = "Project" +def spack_fork = "C2SM" +def spack_branch = "main" +String[] triggerPhrase = env.ghprbCommentBody.split(" ") + +def parseTriggerPhrase(it) { + fork = it.split("=")[1].split("/")[0] + branch = it.split("=")[1].split("/")[1] + return [fork, branch] +} +triggerPhrase.each { + if(it.contains("spack${repo_identifier}")) { + (spack_fork, spack_branch) = parseTriggerPhrase(it) + } +} pipeline { agent none stages { @@ -20,7 +35,7 @@ pipeline { stage('Clone Repos') { steps { sh """ - git clone --depth 1 --recurse-submodules --shallow-submodules https://github.com/C2SM/spack-c2sm.git + git clone --depth 1 --recurse-submodules --shallow-submodules -b '${spack_branch}' https://github.com/${spack_fork}/spack-c2sm.git git clone . local_copy """ } diff --git a/jenkins/spack-PR-icon b/jenkins/spack-PR-icon index 7cadc494df..4d0fe56c22 100644 --- a/jenkins/spack-PR-icon +++ b/jenkins/spack-PR-icon @@ -1,11 +1,9 @@ def repo_identifier = "Project" def icon_fork = "C2SM" def icon_branch = "icon-dsl" -def icon_default = true def spack_fork = "C2SM" def spack_branch = "main" -def spack_default = true def gt4py_fork = "GridTools" def gt4py_branch = "icon4py_20230621" @@ -20,19 +18,16 @@ def parseTriggerPhrase(it) { triggerPhrase.each { if(it.contains("icon${repo_identifier}")) { (icon_fork, icon_branch) = parseTriggerPhrase(it) - icon_default = false } } triggerPhrase.each { if(it.contains("spack${repo_identifier}")) { (spack_fork, spack_branch) = parseTriggerPhrase(it) - spack_default = false } } triggerPhrase.each { if(it.contains("gt4py${repo_identifier}")) { (gt4py_fork, gt4py_branch) = parseTriggerPhrase(it) - gt4py_default = false } } diff --git a/jenkins/spack-PR-stable b/jenkins/spack-PR-stable index 78b85b9961..b2a317a4ff 100644 --- a/jenkins/spack-PR-stable +++ b/jenkins/spack-PR-stable @@ -1,3 +1,18 @@ +def repo_identifier = "Project" +def spack_fork = "C2SM" +def spack_branch = "main" +String[] triggerPhrase = env.ghprbCommentBody.split(" ") + +def parseTriggerPhrase(it) { + fork = it.split("=")[1].split("/")[0] + branch = it.split("=")[1].split("/")[1] + return [fork, branch] +} +triggerPhrase.each { + if(it.contains("spack${repo_identifier}")) { + (spack_fork, spack_branch) = parseTriggerPhrase(it) + } +} pipeline { agent none stages { @@ -20,7 +35,7 @@ pipeline { stage('Clone Repos') { steps { sh """ - git clone --depth 1 --recurse-submodules --shallow-submodules https://github.com/C2SM/spack-c2sm.git + git clone --depth 1 --recurse-submodules --shallow-submodules -b '${spack_branch}' https://github.com/${spack_fork}/spack-c2sm.git git clone . local_copy """ } From 7876eb55714a757c07e6a3e59d31ed4bf71e27df Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 21 Sep 2023 08:23:18 +0200 Subject: [PATCH 097/105] Adapt conftest.py and remove stencil not in code path --- .../test_recon_lsq_cell_c_stencil.py | 657 ------------------ .../advection/recon_lsq_cell_c_stencil.py | 438 ------------ 2 files changed, 1095 deletions(-) delete mode 100644 model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py delete mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py diff --git a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py deleted file mode 100644 index 41ea35df79..0000000000 --- a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_stencil.py +++ /dev/null @@ -1,657 +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 - -import sys # Increase recusion depth, otherwise it doesn't compile - -import numpy as np -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider - -from icon4py.model.atmosphere.advection.recon_lsq_cell_c_stencil import recon_lsq_cell_c_stencil -from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh - - -sys.setrecursionlimit(6000) - - -def recon_lsq_cell_c_stencil_numpy( - c2e2c2e2c: np.ndarray, - p_cc: np.ndarray, - lsq_qtmat_c_1: np.ndarray, - lsq_qtmat_c_2: np.ndarray, - lsq_qtmat_c_3: np.ndarray, - lsq_qtmat_c_4: np.ndarray, - lsq_qtmat_c_5: np.ndarray, - lsq_qtmat_c_6: np.ndarray, - lsq_qtmat_c_7: np.ndarray, - lsq_qtmat_c_8: np.ndarray, - lsq_qtmat_c_9: np.ndarray, - lsq_rmat_rdiag_c_1: np.ndarray, - lsq_rmat_rdiag_c_2: np.ndarray, - lsq_rmat_rdiag_c_3: np.ndarray, - lsq_rmat_rdiag_c_4: np.ndarray, - lsq_rmat_rdiag_c_5: np.ndarray, - lsq_rmat_rdiag_c_6: np.ndarray, - lsq_rmat_rdiag_c_7: np.ndarray, - lsq_rmat_rdiag_c_8: np.ndarray, - lsq_rmat_rdiag_c_9: np.ndarray, - lsq_rmat_utri_c_1: np.ndarray, - lsq_rmat_utri_c_2: np.ndarray, - lsq_rmat_utri_c_3: np.ndarray, - lsq_rmat_utri_c_4: np.ndarray, - lsq_rmat_utri_c_5: np.ndarray, - lsq_rmat_utri_c_6: np.ndarray, - lsq_rmat_utri_c_7: np.ndarray, - lsq_rmat_utri_c_8: np.ndarray, - lsq_rmat_utri_c_9: np.ndarray, - lsq_rmat_utri_c_10: np.ndarray, - lsq_rmat_utri_c_11: np.ndarray, - lsq_rmat_utri_c_12: np.ndarray, - lsq_rmat_utri_c_13: np.ndarray, - lsq_rmat_utri_c_14: np.ndarray, - lsq_rmat_utri_c_15: np.ndarray, - lsq_rmat_utri_c_16: np.ndarray, - lsq_rmat_utri_c_17: np.ndarray, - lsq_rmat_utri_c_18: np.ndarray, - lsq_rmat_utri_c_19: np.ndarray, - lsq_rmat_utri_c_20: np.ndarray, - lsq_rmat_utri_c_21: np.ndarray, - lsq_rmat_utri_c_22: np.ndarray, - lsq_rmat_utri_c_23: np.ndarray, - lsq_rmat_utri_c_24: np.ndarray, - lsq_rmat_utri_c_25: np.ndarray, - lsq_rmat_utri_c_26: np.ndarray, - lsq_rmat_utri_c_27: np.ndarray, - lsq_rmat_utri_c_28: np.ndarray, - lsq_rmat_utri_c_29: np.ndarray, - lsq_rmat_utri_c_30: np.ndarray, - lsq_rmat_utri_c_31: np.ndarray, - lsq_rmat_utri_c_32: np.ndarray, - lsq_rmat_utri_c_33: np.ndarray, - lsq_rmat_utri_c_34: np.ndarray, - lsq_rmat_utri_c_35: np.ndarray, - lsq_rmat_utri_c_36: np.ndarray, - lsq_moments_1: np.ndarray, - lsq_moments_2: np.ndarray, - lsq_moments_3: np.ndarray, - lsq_moments_4: np.ndarray, - lsq_moments_5: np.ndarray, - lsq_moments_6: np.ndarray, - lsq_moments_7: np.ndarray, - lsq_moments_8: np.ndarray, - lsq_moments_9: np.ndarray, -) -> tuple[np.ndarray]: - lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) - lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) - lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) - lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) - lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) - lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) - lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) - lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) - lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) - lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) - lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) - lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) - lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) - lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) - lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) - lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) - lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) - lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) - lsq_rmat_rdiag_c_1 = np.expand_dims(lsq_rmat_rdiag_c_1, axis=-1) - lsq_rmat_rdiag_c_2 = np.expand_dims(lsq_rmat_rdiag_c_2, axis=-1) - lsq_rmat_rdiag_c_3 = np.expand_dims(lsq_rmat_rdiag_c_3, axis=-1) - lsq_rmat_rdiag_c_4 = np.expand_dims(lsq_rmat_rdiag_c_4, axis=-1) - lsq_rmat_rdiag_c_5 = np.expand_dims(lsq_rmat_rdiag_c_5, axis=-1) - lsq_rmat_rdiag_c_6 = np.expand_dims(lsq_rmat_rdiag_c_6, axis=-1) - lsq_rmat_rdiag_c_7 = np.expand_dims(lsq_rmat_rdiag_c_7, axis=-1) - lsq_rmat_rdiag_c_8 = np.expand_dims(lsq_rmat_rdiag_c_8, axis=-1) - lsq_rmat_rdiag_c_9 = np.expand_dims(lsq_rmat_rdiag_c_9, axis=-1) - lsq_rmat_rdiag_c_1 = np.broadcast_to(lsq_rmat_rdiag_c_1, p_cc.shape) - lsq_rmat_rdiag_c_2 = np.broadcast_to(lsq_rmat_rdiag_c_2, p_cc.shape) - lsq_rmat_rdiag_c_3 = np.broadcast_to(lsq_rmat_rdiag_c_3, p_cc.shape) - lsq_rmat_rdiag_c_4 = np.broadcast_to(lsq_rmat_rdiag_c_4, p_cc.shape) - lsq_rmat_rdiag_c_5 = np.broadcast_to(lsq_rmat_rdiag_c_5, p_cc.shape) - lsq_rmat_rdiag_c_6 = np.broadcast_to(lsq_rmat_rdiag_c_6, p_cc.shape) - lsq_rmat_rdiag_c_7 = np.broadcast_to(lsq_rmat_rdiag_c_7, p_cc.shape) - lsq_rmat_rdiag_c_8 = np.broadcast_to(lsq_rmat_rdiag_c_8, p_cc.shape) - lsq_rmat_rdiag_c_9 = np.broadcast_to(lsq_rmat_rdiag_c_9, p_cc.shape) - lsq_rmat_utri_c_1 = np.expand_dims(lsq_rmat_utri_c_1, axis=-1) - lsq_rmat_utri_c_2 = np.expand_dims(lsq_rmat_utri_c_2, axis=-1) - lsq_rmat_utri_c_3 = np.expand_dims(lsq_rmat_utri_c_3, axis=-1) - lsq_rmat_utri_c_4 = np.expand_dims(lsq_rmat_utri_c_4, axis=-1) - lsq_rmat_utri_c_5 = np.expand_dims(lsq_rmat_utri_c_5, axis=-1) - lsq_rmat_utri_c_6 = np.expand_dims(lsq_rmat_utri_c_6, axis=-1) - lsq_rmat_utri_c_7 = np.expand_dims(lsq_rmat_utri_c_7, axis=-1) - lsq_rmat_utri_c_8 = np.expand_dims(lsq_rmat_utri_c_8, axis=-1) - lsq_rmat_utri_c_9 = np.expand_dims(lsq_rmat_utri_c_9, axis=-1) - lsq_rmat_utri_c_10 = np.expand_dims(lsq_rmat_utri_c_10, axis=-1) - lsq_rmat_utri_c_11 = np.expand_dims(lsq_rmat_utri_c_11, axis=-1) - lsq_rmat_utri_c_12 = np.expand_dims(lsq_rmat_utri_c_12, axis=-1) - lsq_rmat_utri_c_13 = np.expand_dims(lsq_rmat_utri_c_13, axis=-1) - lsq_rmat_utri_c_14 = np.expand_dims(lsq_rmat_utri_c_14, axis=-1) - lsq_rmat_utri_c_15 = np.expand_dims(lsq_rmat_utri_c_15, axis=-1) - lsq_rmat_utri_c_16 = np.expand_dims(lsq_rmat_utri_c_16, axis=-1) - lsq_rmat_utri_c_17 = np.expand_dims(lsq_rmat_utri_c_17, axis=-1) - lsq_rmat_utri_c_18 = np.expand_dims(lsq_rmat_utri_c_18, axis=-1) - lsq_rmat_utri_c_19 = np.expand_dims(lsq_rmat_utri_c_19, axis=-1) - lsq_rmat_utri_c_20 = np.expand_dims(lsq_rmat_utri_c_20, axis=-1) - lsq_rmat_utri_c_21 = np.expand_dims(lsq_rmat_utri_c_21, axis=-1) - lsq_rmat_utri_c_22 = np.expand_dims(lsq_rmat_utri_c_22, axis=-1) - lsq_rmat_utri_c_23 = np.expand_dims(lsq_rmat_utri_c_23, axis=-1) - lsq_rmat_utri_c_24 = np.expand_dims(lsq_rmat_utri_c_24, axis=-1) - lsq_rmat_utri_c_25 = np.expand_dims(lsq_rmat_utri_c_25, axis=-1) - lsq_rmat_utri_c_26 = np.expand_dims(lsq_rmat_utri_c_26, axis=-1) - lsq_rmat_utri_c_27 = np.expand_dims(lsq_rmat_utri_c_27, axis=-1) - lsq_rmat_utri_c_28 = np.expand_dims(lsq_rmat_utri_c_28, axis=-1) - lsq_rmat_utri_c_29 = np.expand_dims(lsq_rmat_utri_c_29, axis=-1) - lsq_rmat_utri_c_30 = np.expand_dims(lsq_rmat_utri_c_30, axis=-1) - lsq_rmat_utri_c_31 = np.expand_dims(lsq_rmat_utri_c_31, axis=-1) - lsq_rmat_utri_c_32 = np.expand_dims(lsq_rmat_utri_c_32, axis=-1) - lsq_rmat_utri_c_33 = np.expand_dims(lsq_rmat_utri_c_33, axis=-1) - lsq_rmat_utri_c_34 = np.expand_dims(lsq_rmat_utri_c_34, axis=-1) - lsq_rmat_utri_c_35 = np.expand_dims(lsq_rmat_utri_c_35, axis=-1) - lsq_rmat_utri_c_36 = np.expand_dims(lsq_rmat_utri_c_36, axis=-1) - lsq_rmat_utri_c_1 = np.broadcast_to(lsq_rmat_utri_c_1, p_cc.shape) - lsq_rmat_utri_c_2 = np.broadcast_to(lsq_rmat_utri_c_2, p_cc.shape) - lsq_rmat_utri_c_3 = np.broadcast_to(lsq_rmat_utri_c_3, p_cc.shape) - lsq_rmat_utri_c_4 = np.broadcast_to(lsq_rmat_utri_c_4, p_cc.shape) - lsq_rmat_utri_c_5 = np.broadcast_to(lsq_rmat_utri_c_5, p_cc.shape) - lsq_rmat_utri_c_6 = np.broadcast_to(lsq_rmat_utri_c_6, p_cc.shape) - lsq_rmat_utri_c_7 = np.broadcast_to(lsq_rmat_utri_c_7, p_cc.shape) - lsq_rmat_utri_c_8 = np.broadcast_to(lsq_rmat_utri_c_8, p_cc.shape) - lsq_rmat_utri_c_9 = np.broadcast_to(lsq_rmat_utri_c_9, p_cc.shape) - lsq_rmat_utri_c_10 = np.broadcast_to(lsq_rmat_utri_c_10, p_cc.shape) - lsq_rmat_utri_c_11 = np.broadcast_to(lsq_rmat_utri_c_11, p_cc.shape) - lsq_rmat_utri_c_12 = np.broadcast_to(lsq_rmat_utri_c_12, p_cc.shape) - lsq_rmat_utri_c_13 = np.broadcast_to(lsq_rmat_utri_c_13, p_cc.shape) - lsq_rmat_utri_c_14 = np.broadcast_to(lsq_rmat_utri_c_14, p_cc.shape) - lsq_rmat_utri_c_15 = np.broadcast_to(lsq_rmat_utri_c_15, p_cc.shape) - lsq_rmat_utri_c_16 = np.broadcast_to(lsq_rmat_utri_c_16, p_cc.shape) - lsq_rmat_utri_c_17 = np.broadcast_to(lsq_rmat_utri_c_17, p_cc.shape) - lsq_rmat_utri_c_18 = np.broadcast_to(lsq_rmat_utri_c_18, p_cc.shape) - lsq_rmat_utri_c_19 = np.broadcast_to(lsq_rmat_utri_c_19, p_cc.shape) - lsq_rmat_utri_c_20 = np.broadcast_to(lsq_rmat_utri_c_20, p_cc.shape) - lsq_rmat_utri_c_21 = np.broadcast_to(lsq_rmat_utri_c_21, p_cc.shape) - lsq_rmat_utri_c_22 = np.broadcast_to(lsq_rmat_utri_c_22, p_cc.shape) - lsq_rmat_utri_c_23 = np.broadcast_to(lsq_rmat_utri_c_23, p_cc.shape) - lsq_rmat_utri_c_24 = np.broadcast_to(lsq_rmat_utri_c_24, p_cc.shape) - lsq_rmat_utri_c_25 = np.broadcast_to(lsq_rmat_utri_c_25, p_cc.shape) - lsq_rmat_utri_c_26 = np.broadcast_to(lsq_rmat_utri_c_26, p_cc.shape) - lsq_rmat_utri_c_27 = np.broadcast_to(lsq_rmat_utri_c_27, p_cc.shape) - lsq_rmat_utri_c_28 = np.broadcast_to(lsq_rmat_utri_c_28, p_cc.shape) - lsq_rmat_utri_c_29 = np.broadcast_to(lsq_rmat_utri_c_29, p_cc.shape) - lsq_rmat_utri_c_30 = np.broadcast_to(lsq_rmat_utri_c_30, p_cc.shape) - lsq_rmat_utri_c_31 = np.broadcast_to(lsq_rmat_utri_c_31, p_cc.shape) - lsq_rmat_utri_c_32 = np.broadcast_to(lsq_rmat_utri_c_32, p_cc.shape) - lsq_rmat_utri_c_33 = np.broadcast_to(lsq_rmat_utri_c_33, p_cc.shape) - lsq_rmat_utri_c_34 = np.broadcast_to(lsq_rmat_utri_c_34, p_cc.shape) - lsq_rmat_utri_c_35 = np.broadcast_to(lsq_rmat_utri_c_35, p_cc.shape) - lsq_rmat_utri_c_36 = np.broadcast_to(lsq_rmat_utri_c_36, p_cc.shape) - lsq_qtmat_c_1 = np.expand_dims(lsq_qtmat_c_1, axis=-1) - lsq_qtmat_c_2 = np.expand_dims(lsq_qtmat_c_2, axis=-1) - lsq_qtmat_c_3 = np.expand_dims(lsq_qtmat_c_3, axis=-1) - lsq_qtmat_c_4 = np.expand_dims(lsq_qtmat_c_4, axis=-1) - lsq_qtmat_c_5 = np.expand_dims(lsq_qtmat_c_5, axis=-1) - lsq_qtmat_c_6 = np.expand_dims(lsq_qtmat_c_6, axis=-1) - lsq_qtmat_c_7 = np.expand_dims(lsq_qtmat_c_7, axis=-1) - lsq_qtmat_c_8 = np.expand_dims(lsq_qtmat_c_8, axis=-1) - lsq_qtmat_c_9 = np.expand_dims(lsq_qtmat_c_9, axis=-1) - - p_coeff_10 = lsq_rmat_rdiag_c_9 * ( - lsq_qtmat_c_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_9 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - lsq_rmat_utri_c_1 * p_coeff_10 - ) - - p_coeff_8 = lsq_rmat_rdiag_c_7 * ( - lsq_qtmat_c_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) - ) - - p_coeff_7 = lsq_rmat_rdiag_c_6 * ( - lsq_qtmat_c_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_4 * p_coeff_8 - + lsq_rmat_utri_c_5 * p_coeff_9 - + lsq_rmat_utri_c_6 * p_coeff_10 - ) - ) - - p_coeff_6 = lsq_rmat_rdiag_c_5 * ( - lsq_qtmat_c_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_7 * p_coeff_7 - + lsq_rmat_utri_c_8 * p_coeff_8 - + lsq_rmat_utri_c_9 * p_coeff_9 - + lsq_rmat_utri_c_10 * p_coeff_10 - ) - ) - - p_coeff_5 = lsq_rmat_rdiag_c_4 * ( - lsq_qtmat_c_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_11 * p_coeff_6 - + lsq_rmat_utri_c_12 * p_coeff_7 - + lsq_rmat_utri_c_13 * p_coeff_8 - + lsq_rmat_utri_c_14 * p_coeff_9 - + lsq_rmat_utri_c_15 * p_coeff_10 - ) - ) - - p_coeff_4 = lsq_rmat_rdiag_c_3 * ( - lsq_qtmat_c_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_16 * p_coeff_5 - + lsq_rmat_utri_c_17 * p_coeff_6 - + lsq_rmat_utri_c_18 * p_coeff_7 - + lsq_rmat_utri_c_19 * p_coeff_8 - + lsq_rmat_utri_c_20 * p_coeff_9 - + lsq_rmat_utri_c_21 * p_coeff_10 - ) - ) - - p_coeff_3 = lsq_rmat_rdiag_c_2 * ( - lsq_qtmat_c_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_22 * p_coeff_4 - + lsq_rmat_utri_c_23 * p_coeff_5 - + lsq_rmat_utri_c_24 * p_coeff_6 - + lsq_rmat_utri_c_25 * p_coeff_7 - + lsq_rmat_utri_c_26 * p_coeff_8 - + lsq_rmat_utri_c_27 * p_coeff_9 - + lsq_rmat_utri_c_28 * p_coeff_10 - ) - ) - - p_coeff_2 = lsq_rmat_rdiag_c_1 * ( - lsq_qtmat_c_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_qtmat_c_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_qtmat_c_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_qtmat_c_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_qtmat_c_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_qtmat_c_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_qtmat_c_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_qtmat_c_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_qtmat_c_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - - ( - lsq_rmat_utri_c_29 * p_coeff_3 - + lsq_rmat_utri_c_30 * p_coeff_4 - + lsq_rmat_utri_c_31 * p_coeff_5 - + lsq_rmat_utri_c_32 * p_coeff_6 - + lsq_rmat_utri_c_33 * p_coeff_7 - + lsq_rmat_utri_c_34 * p_coeff_8 - + lsq_rmat_utri_c_35 * p_coeff_9 - + lsq_rmat_utri_c_36 * p_coeff_10 - ) - ) - - p_coeff_1 = p_cc - ( - p_coeff_2 * lsq_moments_1 - + p_coeff_3 * lsq_moments_2 - + p_coeff_4 * lsq_moments_3 - + p_coeff_5 * lsq_moments_4 - + p_coeff_6 * lsq_moments_5 - + p_coeff_7 * lsq_moments_6 - + p_coeff_8 * lsq_moments_7 - + p_coeff_9 * lsq_moments_8 - + p_coeff_10 * lsq_moments_9 - ) - return ( - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, - ) - - -def test_recon_lsq_cell_c_stencil(): - mesh = SimpleMesh() - p_cc = random_field(mesh, CellDim, KDim) - lsq_qtmat_c_1 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_2 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_3 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_4 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_5 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_6 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_7 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_8 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_9 = random_field(mesh, CellDim, C2E2C2E2CDim) - lsq_qtmat_c_1_field = as_1D_sparse_field(lsq_qtmat_c_1, CECECDim) - lsq_qtmat_c_2_field = as_1D_sparse_field(lsq_qtmat_c_2, CECECDim) - lsq_qtmat_c_3_field = as_1D_sparse_field(lsq_qtmat_c_3, CECECDim) - lsq_qtmat_c_4_field = as_1D_sparse_field(lsq_qtmat_c_4, CECECDim) - lsq_qtmat_c_5_field = as_1D_sparse_field(lsq_qtmat_c_5, CECECDim) - lsq_qtmat_c_6_field = as_1D_sparse_field(lsq_qtmat_c_6, CECECDim) - lsq_qtmat_c_7_field = as_1D_sparse_field(lsq_qtmat_c_7, CECECDim) - lsq_qtmat_c_8_field = as_1D_sparse_field(lsq_qtmat_c_8, CECECDim) - lsq_qtmat_c_9_field = as_1D_sparse_field(lsq_qtmat_c_9, CECECDim) - lsq_rmat_rdiag_c_1 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_2 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_3 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_4 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_5 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_6 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_7 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_8 = random_field(mesh, CellDim) - lsq_rmat_rdiag_c_9 = random_field(mesh, CellDim) - lsq_rmat_utri_c_1 = random_field(mesh, CellDim) - lsq_rmat_utri_c_2 = random_field(mesh, CellDim) - lsq_rmat_utri_c_3 = random_field(mesh, CellDim) - lsq_rmat_utri_c_4 = random_field(mesh, CellDim) - lsq_rmat_utri_c_5 = random_field(mesh, CellDim) - lsq_rmat_utri_c_6 = random_field(mesh, CellDim) - lsq_rmat_utri_c_7 = random_field(mesh, CellDim) - lsq_rmat_utri_c_8 = random_field(mesh, CellDim) - lsq_rmat_utri_c_9 = random_field(mesh, CellDim) - lsq_rmat_utri_c_10 = random_field(mesh, CellDim) - lsq_rmat_utri_c_11 = random_field(mesh, CellDim) - lsq_rmat_utri_c_12 = random_field(mesh, CellDim) - lsq_rmat_utri_c_13 = random_field(mesh, CellDim) - lsq_rmat_utri_c_14 = random_field(mesh, CellDim) - lsq_rmat_utri_c_15 = random_field(mesh, CellDim) - lsq_rmat_utri_c_16 = random_field(mesh, CellDim) - lsq_rmat_utri_c_17 = random_field(mesh, CellDim) - lsq_rmat_utri_c_18 = random_field(mesh, CellDim) - lsq_rmat_utri_c_19 = random_field(mesh, CellDim) - lsq_rmat_utri_c_20 = random_field(mesh, CellDim) - lsq_rmat_utri_c_21 = random_field(mesh, CellDim) - lsq_rmat_utri_c_22 = random_field(mesh, CellDim) - lsq_rmat_utri_c_23 = random_field(mesh, CellDim) - lsq_rmat_utri_c_24 = random_field(mesh, CellDim) - lsq_rmat_utri_c_25 = random_field(mesh, CellDim) - lsq_rmat_utri_c_26 = random_field(mesh, CellDim) - lsq_rmat_utri_c_27 = random_field(mesh, CellDim) - lsq_rmat_utri_c_28 = random_field(mesh, CellDim) - lsq_rmat_utri_c_29 = random_field(mesh, CellDim) - lsq_rmat_utri_c_30 = random_field(mesh, CellDim) - lsq_rmat_utri_c_31 = random_field(mesh, CellDim) - lsq_rmat_utri_c_32 = random_field(mesh, CellDim) - lsq_rmat_utri_c_33 = random_field(mesh, CellDim) - lsq_rmat_utri_c_34 = random_field(mesh, CellDim) - lsq_rmat_utri_c_35 = random_field(mesh, CellDim) - lsq_rmat_utri_c_36 = random_field(mesh, CellDim) - lsq_moments_1 = random_field(mesh, CellDim) - lsq_moments_2 = random_field(mesh, CellDim) - lsq_moments_3 = random_field(mesh, CellDim) - lsq_moments_4 = random_field(mesh, CellDim) - lsq_moments_5 = random_field(mesh, CellDim) - lsq_moments_6 = random_field(mesh, CellDim) - lsq_moments_7 = random_field(mesh, CellDim) - lsq_moments_8 = random_field(mesh, CellDim) - lsq_moments_9 = random_field(mesh, CellDim) - p_coeff_1 = zero_field(mesh, CellDim, KDim) - p_coeff_2 = zero_field(mesh, CellDim, KDim) - p_coeff_3 = zero_field(mesh, CellDim, KDim) - p_coeff_4 = zero_field(mesh, CellDim, KDim) - p_coeff_5 = zero_field(mesh, CellDim, KDim) - p_coeff_6 = zero_field(mesh, CellDim, KDim) - p_coeff_7 = zero_field(mesh, CellDim, KDim) - p_coeff_8 = zero_field(mesh, CellDim, KDim) - p_coeff_9 = zero_field(mesh, CellDim, KDim) - p_coeff_10 = zero_field(mesh, CellDim, KDim) - - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ) = recon_lsq_cell_c_stencil_numpy( - mesh.c2e2c2e2c, - np.asarray(p_cc), - np.asarray(lsq_qtmat_c_1), - np.asarray(lsq_qtmat_c_2), - np.asarray(lsq_qtmat_c_3), - np.asarray(lsq_qtmat_c_4), - np.asarray(lsq_qtmat_c_5), - np.asarray(lsq_qtmat_c_6), - np.asarray(lsq_qtmat_c_7), - np.asarray(lsq_qtmat_c_8), - np.asarray(lsq_qtmat_c_9), - np.asarray(lsq_rmat_rdiag_c_1), - np.asarray(lsq_rmat_rdiag_c_2), - np.asarray(lsq_rmat_rdiag_c_3), - np.asarray(lsq_rmat_rdiag_c_4), - np.asarray(lsq_rmat_rdiag_c_5), - np.asarray(lsq_rmat_rdiag_c_6), - np.asarray(lsq_rmat_rdiag_c_7), - np.asarray(lsq_rmat_rdiag_c_8), - np.asarray(lsq_rmat_rdiag_c_9), - np.asarray(lsq_rmat_utri_c_1), - np.asarray(lsq_rmat_utri_c_2), - np.asarray(lsq_rmat_utri_c_3), - np.asarray(lsq_rmat_utri_c_4), - np.asarray(lsq_rmat_utri_c_5), - np.asarray(lsq_rmat_utri_c_6), - np.asarray(lsq_rmat_utri_c_7), - np.asarray(lsq_rmat_utri_c_8), - np.asarray(lsq_rmat_utri_c_9), - np.asarray(lsq_rmat_utri_c_10), - np.asarray(lsq_rmat_utri_c_11), - np.asarray(lsq_rmat_utri_c_12), - np.asarray(lsq_rmat_utri_c_13), - np.asarray(lsq_rmat_utri_c_14), - np.asarray(lsq_rmat_utri_c_15), - np.asarray(lsq_rmat_utri_c_16), - np.asarray(lsq_rmat_utri_c_17), - np.asarray(lsq_rmat_utri_c_18), - np.asarray(lsq_rmat_utri_c_19), - np.asarray(lsq_rmat_utri_c_20), - np.asarray(lsq_rmat_utri_c_21), - np.asarray(lsq_rmat_utri_c_22), - np.asarray(lsq_rmat_utri_c_23), - np.asarray(lsq_rmat_utri_c_24), - np.asarray(lsq_rmat_utri_c_25), - np.asarray(lsq_rmat_utri_c_26), - np.asarray(lsq_rmat_utri_c_27), - np.asarray(lsq_rmat_utri_c_28), - np.asarray(lsq_rmat_utri_c_29), - np.asarray(lsq_rmat_utri_c_30), - np.asarray(lsq_rmat_utri_c_31), - np.asarray(lsq_rmat_utri_c_32), - np.asarray(lsq_rmat_utri_c_33), - np.asarray(lsq_rmat_utri_c_34), - np.asarray(lsq_rmat_utri_c_35), - np.asarray(lsq_rmat_utri_c_36), - np.asarray(lsq_moments_1), - np.asarray(lsq_moments_2), - np.asarray(lsq_moments_3), - np.asarray(lsq_moments_4), - np.asarray(lsq_moments_5), - np.asarray(lsq_moments_6), - np.asarray(lsq_moments_7), - np.asarray(lsq_moments_8), - np.asarray(lsq_moments_9), - ) - - recon_lsq_cell_c_stencil( - p_cc, - lsq_qtmat_c_1_field, - lsq_qtmat_c_2_field, - lsq_qtmat_c_3_field, - lsq_qtmat_c_4_field, - lsq_qtmat_c_5_field, - lsq_qtmat_c_6_field, - lsq_qtmat_c_7_field, - lsq_qtmat_c_8_field, - lsq_qtmat_c_9_field, - lsq_rmat_rdiag_c_1, - lsq_rmat_rdiag_c_2, - lsq_rmat_rdiag_c_3, - lsq_rmat_rdiag_c_4, - lsq_rmat_rdiag_c_5, - lsq_rmat_rdiag_c_6, - lsq_rmat_rdiag_c_7, - lsq_rmat_rdiag_c_8, - lsq_rmat_rdiag_c_9, - lsq_rmat_utri_c_1, - lsq_rmat_utri_c_2, - lsq_rmat_utri_c_3, - lsq_rmat_utri_c_4, - lsq_rmat_utri_c_5, - lsq_rmat_utri_c_6, - lsq_rmat_utri_c_7, - lsq_rmat_utri_c_8, - lsq_rmat_utri_c_9, - lsq_rmat_utri_c_10, - lsq_rmat_utri_c_11, - lsq_rmat_utri_c_12, - lsq_rmat_utri_c_13, - lsq_rmat_utri_c_14, - lsq_rmat_utri_c_15, - lsq_rmat_utri_c_16, - lsq_rmat_utri_c_17, - lsq_rmat_utri_c_18, - lsq_rmat_utri_c_19, - lsq_rmat_utri_c_20, - lsq_rmat_utri_c_21, - lsq_rmat_utri_c_22, - lsq_rmat_utri_c_23, - lsq_rmat_utri_c_24, - lsq_rmat_utri_c_25, - lsq_rmat_utri_c_26, - lsq_rmat_utri_c_27, - lsq_rmat_utri_c_28, - lsq_rmat_utri_c_29, - lsq_rmat_utri_c_30, - lsq_rmat_utri_c_31, - lsq_rmat_utri_c_32, - lsq_rmat_utri_c_33, - lsq_rmat_utri_c_34, - lsq_rmat_utri_c_35, - lsq_rmat_utri_c_36, - lsq_moments_1, - lsq_moments_2, - lsq_moments_3, - lsq_moments_4, - lsq_moments_5, - lsq_moments_6, - lsq_moments_7, - lsq_moments_8, - lsq_moments_9, - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, - offset_provider={ - "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), - "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), - }, - ) - co1 = np.asarray(p_coeff_1) - co2 = np.asarray(p_coeff_2) - co3 = np.asarray(p_coeff_3) - co4 = np.asarray(p_coeff_4) - co5 = np.asarray(p_coeff_5) - co6 = np.asarray(p_coeff_6) - co7 = np.asarray(p_coeff_7) - co8 = np.asarray(p_coeff_8) - co9 = np.asarray(p_coeff_9) - co10 = np.asarray(p_coeff_10) - assert np.allclose(ref_1, co1) - assert np.allclose(ref_2, co2) - assert np.allclose(ref_3, co3) - assert np.allclose(ref_4, co4) - assert np.allclose(ref_5, co5) - assert np.allclose(ref_6, co6) - assert np.allclose(ref_7, co7) - assert np.allclose(ref_8, co8) - assert np.allclose(ref_9, co9) - assert np.allclose(ref_10, co10) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py deleted file mode 100644 index 4866b8f2dd..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_stencil.py +++ /dev/null @@ -1,438 +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 - -import sys # Increase recusion depth, otherwise it doesn't compile - -from gt4py.next.common import Field -from gt4py.next.ffront.decorator import field_operator, program - -from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim - - -sys.setrecursionlimit(6000) - - -@field_operator -def _recon_lsq_cell_c_stencil( - p_cc: Field[[CellDim, KDim], float], - lsq_qtmat_c_1: Field[[CECECDim], float], - lsq_qtmat_c_2: Field[[CECECDim], float], - lsq_qtmat_c_3: Field[[CECECDim], float], - lsq_qtmat_c_4: Field[[CECECDim], float], - lsq_qtmat_c_5: Field[[CECECDim], float], - lsq_qtmat_c_6: Field[[CECECDim], float], - lsq_qtmat_c_7: Field[[CECECDim], float], - lsq_qtmat_c_8: Field[[CECECDim], float], - lsq_qtmat_c_9: Field[[CECECDim], float], - lsq_rmat_rdiag_c_1: Field[[CellDim], float], - lsq_rmat_rdiag_c_2: Field[[CellDim], float], - lsq_rmat_rdiag_c_3: Field[[CellDim], float], - lsq_rmat_rdiag_c_4: Field[[CellDim], float], - lsq_rmat_rdiag_c_5: Field[[CellDim], float], - lsq_rmat_rdiag_c_6: Field[[CellDim], float], - lsq_rmat_rdiag_c_7: Field[[CellDim], float], - lsq_rmat_rdiag_c_8: Field[[CellDim], float], - lsq_rmat_rdiag_c_9: Field[[CellDim], float], - lsq_rmat_utri_c_1: Field[[CellDim], float], - lsq_rmat_utri_c_2: Field[[CellDim], float], - lsq_rmat_utri_c_3: Field[[CellDim], float], - lsq_rmat_utri_c_4: Field[[CellDim], float], - lsq_rmat_utri_c_5: Field[[CellDim], float], - lsq_rmat_utri_c_6: Field[[CellDim], float], - lsq_rmat_utri_c_7: Field[[CellDim], float], - lsq_rmat_utri_c_8: Field[[CellDim], float], - lsq_rmat_utri_c_9: Field[[CellDim], float], - lsq_rmat_utri_c_10: Field[[CellDim], float], - lsq_rmat_utri_c_11: Field[[CellDim], float], - lsq_rmat_utri_c_12: Field[[CellDim], float], - lsq_rmat_utri_c_13: Field[[CellDim], float], - lsq_rmat_utri_c_14: Field[[CellDim], float], - lsq_rmat_utri_c_15: Field[[CellDim], float], - lsq_rmat_utri_c_16: Field[[CellDim], float], - lsq_rmat_utri_c_17: Field[[CellDim], float], - lsq_rmat_utri_c_18: Field[[CellDim], float], - lsq_rmat_utri_c_19: Field[[CellDim], float], - lsq_rmat_utri_c_20: Field[[CellDim], float], - lsq_rmat_utri_c_21: Field[[CellDim], float], - lsq_rmat_utri_c_22: Field[[CellDim], float], - lsq_rmat_utri_c_23: Field[[CellDim], float], - lsq_rmat_utri_c_24: Field[[CellDim], float], - lsq_rmat_utri_c_25: Field[[CellDim], float], - lsq_rmat_utri_c_26: Field[[CellDim], float], - lsq_rmat_utri_c_27: Field[[CellDim], float], - lsq_rmat_utri_c_28: Field[[CellDim], float], - lsq_rmat_utri_c_29: Field[[CellDim], float], - lsq_rmat_utri_c_30: Field[[CellDim], float], - lsq_rmat_utri_c_31: Field[[CellDim], float], - lsq_rmat_utri_c_32: Field[[CellDim], float], - lsq_rmat_utri_c_33: Field[[CellDim], float], - lsq_rmat_utri_c_34: Field[[CellDim], float], - lsq_rmat_utri_c_35: Field[[CellDim], float], - lsq_rmat_utri_c_36: Field[[CellDim], float], - lsq_moments_1: Field[[CellDim], float], - lsq_moments_2: Field[[CellDim], float], - lsq_moments_3: Field[[CellDim], float], - lsq_moments_4: Field[[CellDim], float], - lsq_moments_5: Field[[CellDim], float], - lsq_moments_6: Field[[CellDim], float], - lsq_moments_7: Field[[CellDim], float], - lsq_moments_8: Field[[CellDim], float], - lsq_moments_9: Field[[CellDim], float], -) -> tuple[ - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], -]: - - p_coeff_10 = lsq_rmat_rdiag_c_9 * ( - lsq_qtmat_c_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_9(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - ) - - p_coeff_9 = lsq_rmat_rdiag_c_8 * ( - lsq_qtmat_c_8(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_8(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - lsq_rmat_utri_c_1 * p_coeff_10 - ) - - p_coeff_8 = lsq_rmat_rdiag_c_7 * ( - lsq_qtmat_c_7(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_7(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - (lsq_rmat_utri_c_2 * p_coeff_9 + lsq_rmat_utri_c_3 * p_coeff_10) - ) - - p_coeff_7 = lsq_rmat_rdiag_c_6 * ( - lsq_qtmat_c_6(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_6(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_4 * p_coeff_8 - + lsq_rmat_utri_c_5 * p_coeff_9 - + lsq_rmat_utri_c_6 * p_coeff_10 - ) - ) - - p_coeff_6 = lsq_rmat_rdiag_c_5 * ( - lsq_qtmat_c_5(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_5(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_7 * p_coeff_7 - + lsq_rmat_utri_c_8 * p_coeff_8 - + lsq_rmat_utri_c_9 * p_coeff_9 - + lsq_rmat_utri_c_10 * p_coeff_10 - ) - ) - - p_coeff_5 = lsq_rmat_rdiag_c_4 * ( - lsq_qtmat_c_4(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_4(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_11 * p_coeff_6 - + lsq_rmat_utri_c_12 * p_coeff_7 - + lsq_rmat_utri_c_13 * p_coeff_8 - + lsq_rmat_utri_c_14 * p_coeff_9 - + lsq_rmat_utri_c_15 * p_coeff_10 - ) - ) - - p_coeff_4 = lsq_rmat_rdiag_c_3 * ( - lsq_qtmat_c_3(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_3(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_16 * p_coeff_5 - + lsq_rmat_utri_c_17 * p_coeff_6 - + lsq_rmat_utri_c_18 * p_coeff_7 - + lsq_rmat_utri_c_19 * p_coeff_8 - + lsq_rmat_utri_c_20 * p_coeff_9 - + lsq_rmat_utri_c_21 * p_coeff_10 - ) - ) - - p_coeff_3 = lsq_rmat_rdiag_c_2 * ( - lsq_qtmat_c_2(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_2(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_22 * p_coeff_4 - + lsq_rmat_utri_c_23 * p_coeff_5 - + lsq_rmat_utri_c_24 * p_coeff_6 - + lsq_rmat_utri_c_25 * p_coeff_7 - + lsq_rmat_utri_c_26 * p_coeff_8 - + lsq_rmat_utri_c_27 * p_coeff_9 - + lsq_rmat_utri_c_28 * p_coeff_10 - ) - ) - - p_coeff_2 = lsq_rmat_rdiag_c_1 * ( - lsq_qtmat_c_1(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) - + lsq_qtmat_c_1(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) - - ( - lsq_rmat_utri_c_29 * p_coeff_3 - + lsq_rmat_utri_c_30 * p_coeff_4 - + lsq_rmat_utri_c_31 * p_coeff_5 - + lsq_rmat_utri_c_32 * p_coeff_6 - + lsq_rmat_utri_c_33 * p_coeff_7 - + lsq_rmat_utri_c_34 * p_coeff_8 - + lsq_rmat_utri_c_35 * p_coeff_9 - + lsq_rmat_utri_c_36 * p_coeff_10 - ) - ) - - p_coeff_1 = p_cc - ( - p_coeff_2 * lsq_moments_1 - + p_coeff_3 * lsq_moments_2 - + p_coeff_4 * lsq_moments_3 - + p_coeff_5 * lsq_moments_4 - + p_coeff_6 * lsq_moments_5 - + p_coeff_7 * lsq_moments_6 - + p_coeff_8 * lsq_moments_7 - + p_coeff_9 * lsq_moments_8 - + p_coeff_10 * lsq_moments_9 - ) - return ( - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, - ) - - -@program -def recon_lsq_cell_c_stencil( - p_cc: Field[[CellDim, KDim], float], - lsq_qtmat_c_1: Field[[CECECDim], float], - lsq_qtmat_c_2: Field[[CECECDim], float], - lsq_qtmat_c_3: Field[[CECECDim], float], - lsq_qtmat_c_4: Field[[CECECDim], float], - lsq_qtmat_c_5: Field[[CECECDim], float], - lsq_qtmat_c_6: Field[[CECECDim], float], - lsq_qtmat_c_7: Field[[CECECDim], float], - lsq_qtmat_c_8: Field[[CECECDim], float], - lsq_qtmat_c_9: Field[[CECECDim], float], - lsq_rmat_rdiag_c_1: Field[[CellDim], float], - lsq_rmat_rdiag_c_2: Field[[CellDim], float], - lsq_rmat_rdiag_c_3: Field[[CellDim], float], - lsq_rmat_rdiag_c_4: Field[[CellDim], float], - lsq_rmat_rdiag_c_5: Field[[CellDim], float], - lsq_rmat_rdiag_c_6: Field[[CellDim], float], - lsq_rmat_rdiag_c_7: Field[[CellDim], float], - lsq_rmat_rdiag_c_8: Field[[CellDim], float], - lsq_rmat_rdiag_c_9: Field[[CellDim], float], - lsq_rmat_utri_c_1: Field[[CellDim], float], - lsq_rmat_utri_c_2: Field[[CellDim], float], - lsq_rmat_utri_c_3: Field[[CellDim], float], - lsq_rmat_utri_c_4: Field[[CellDim], float], - lsq_rmat_utri_c_5: Field[[CellDim], float], - lsq_rmat_utri_c_6: Field[[CellDim], float], - lsq_rmat_utri_c_7: Field[[CellDim], float], - lsq_rmat_utri_c_8: Field[[CellDim], float], - lsq_rmat_utri_c_9: Field[[CellDim], float], - lsq_rmat_utri_c_10: Field[[CellDim], float], - lsq_rmat_utri_c_11: Field[[CellDim], float], - lsq_rmat_utri_c_12: Field[[CellDim], float], - lsq_rmat_utri_c_13: Field[[CellDim], float], - lsq_rmat_utri_c_14: Field[[CellDim], float], - lsq_rmat_utri_c_15: Field[[CellDim], float], - lsq_rmat_utri_c_16: Field[[CellDim], float], - lsq_rmat_utri_c_17: Field[[CellDim], float], - lsq_rmat_utri_c_18: Field[[CellDim], float], - lsq_rmat_utri_c_19: Field[[CellDim], float], - lsq_rmat_utri_c_20: Field[[CellDim], float], - lsq_rmat_utri_c_21: Field[[CellDim], float], - lsq_rmat_utri_c_22: Field[[CellDim], float], - lsq_rmat_utri_c_23: Field[[CellDim], float], - lsq_rmat_utri_c_24: Field[[CellDim], float], - lsq_rmat_utri_c_25: Field[[CellDim], float], - lsq_rmat_utri_c_26: Field[[CellDim], float], - lsq_rmat_utri_c_27: Field[[CellDim], float], - lsq_rmat_utri_c_28: Field[[CellDim], float], - lsq_rmat_utri_c_29: Field[[CellDim], float], - lsq_rmat_utri_c_30: Field[[CellDim], float], - lsq_rmat_utri_c_31: Field[[CellDim], float], - lsq_rmat_utri_c_32: Field[[CellDim], float], - lsq_rmat_utri_c_33: Field[[CellDim], float], - lsq_rmat_utri_c_34: Field[[CellDim], float], - lsq_rmat_utri_c_35: Field[[CellDim], float], - lsq_rmat_utri_c_36: Field[[CellDim], float], - lsq_moments_1: Field[[CellDim], float], - lsq_moments_2: Field[[CellDim], float], - lsq_moments_3: Field[[CellDim], float], - lsq_moments_4: Field[[CellDim], float], - lsq_moments_5: Field[[CellDim], float], - lsq_moments_6: Field[[CellDim], float], - lsq_moments_7: Field[[CellDim], float], - lsq_moments_8: Field[[CellDim], float], - lsq_moments_9: Field[[CellDim], float], - p_coeff_1: Field[[CellDim, KDim], float], - p_coeff_2: Field[[CellDim, KDim], float], - p_coeff_3: Field[[CellDim, KDim], float], - p_coeff_4: Field[[CellDim, KDim], float], - p_coeff_5: Field[[CellDim, KDim], float], - p_coeff_6: Field[[CellDim, KDim], float], - p_coeff_7: Field[[CellDim, KDim], float], - p_coeff_8: Field[[CellDim, KDim], float], - p_coeff_9: Field[[CellDim, KDim], float], - p_coeff_10: Field[[CellDim, KDim], float], -): - _recon_lsq_cell_c_stencil( - p_cc, - lsq_qtmat_c_1, - lsq_qtmat_c_2, - lsq_qtmat_c_3, - lsq_qtmat_c_4, - lsq_qtmat_c_5, - lsq_qtmat_c_6, - lsq_qtmat_c_7, - lsq_qtmat_c_8, - lsq_qtmat_c_9, - lsq_rmat_rdiag_c_1, - lsq_rmat_rdiag_c_2, - lsq_rmat_rdiag_c_3, - lsq_rmat_rdiag_c_4, - lsq_rmat_rdiag_c_5, - lsq_rmat_rdiag_c_6, - lsq_rmat_rdiag_c_7, - lsq_rmat_rdiag_c_8, - lsq_rmat_rdiag_c_9, - lsq_rmat_utri_c_1, - lsq_rmat_utri_c_2, - lsq_rmat_utri_c_3, - lsq_rmat_utri_c_4, - lsq_rmat_utri_c_5, - lsq_rmat_utri_c_6, - lsq_rmat_utri_c_7, - lsq_rmat_utri_c_8, - lsq_rmat_utri_c_9, - lsq_rmat_utri_c_10, - lsq_rmat_utri_c_11, - lsq_rmat_utri_c_12, - lsq_rmat_utri_c_13, - lsq_rmat_utri_c_14, - lsq_rmat_utri_c_15, - lsq_rmat_utri_c_16, - lsq_rmat_utri_c_17, - lsq_rmat_utri_c_18, - lsq_rmat_utri_c_19, - lsq_rmat_utri_c_20, - lsq_rmat_utri_c_21, - lsq_rmat_utri_c_22, - lsq_rmat_utri_c_23, - lsq_rmat_utri_c_24, - lsq_rmat_utri_c_25, - lsq_rmat_utri_c_26, - lsq_rmat_utri_c_27, - lsq_rmat_utri_c_28, - lsq_rmat_utri_c_29, - lsq_rmat_utri_c_30, - lsq_rmat_utri_c_31, - lsq_rmat_utri_c_32, - lsq_rmat_utri_c_33, - lsq_rmat_utri_c_34, - lsq_rmat_utri_c_35, - lsq_rmat_utri_c_36, - lsq_moments_1, - lsq_moments_2, - lsq_moments_3, - lsq_moments_4, - lsq_moments_5, - lsq_moments_6, - lsq_moments_7, - lsq_moments_8, - lsq_moments_9, - out=( - p_coeff_1, - p_coeff_2, - p_coeff_3, - p_coeff_4, - p_coeff_5, - p_coeff_6, - p_coeff_7, - p_coeff_8, - p_coeff_9, - p_coeff_10, - ), - ) From 6a7d8f1155fb659ebf0b60c4fdd5a86009292c22 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Tue, 19 Sep 2023 11:43:33 +0200 Subject: [PATCH 098/105] Fix test references (assuming wrap around) (#270) This fixes test references where the current implementation assumes a certain behavior of the backend. Specifically, we used np.roll in the references and the backend (by chance) had the same wrap-around behavior. For an upcoming PR in GT4Py, we need to remove that expectation. --------- Co-authored-by: Jonas Jucker Co-authored-by: juckerj <39263956+jonasjucker@users.noreply.github.com> Co-authored-by: Abishek Gopal --- .gitpod.Dockerfile | 2 +- jenkins/spack-PR | 21 +++++++++++++------ jenkins/spack-PR-stable | 21 +++++++++++++------ .../dycore/mo_solve_nonhydro_stencil_05.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_08.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_09.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_10.py | 2 +- .../mo_solve_nonhydro_stencil_11_upper.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_36.py | 4 +++- .../dycore/mo_solve_nonhydro_stencil_38.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_39.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_40.py | 2 +- .../mo_velocity_advection_stencil_02.py | 2 +- .../mo_velocity_advection_stencil_03.py | 2 +- .../mo_velocity_advection_stencil_06.py | 2 +- .../mo_velocity_advection_stencil_10.py | 2 +- .../test_mo_solve_nonhydro_stencil_05.py | 1 + .../test_mo_solve_nonhydro_stencil_08.py | 4 ++++ .../test_mo_solve_nonhydro_stencil_09.py | 3 +++ .../test_mo_solve_nonhydro_stencil_10.py | 14 ++++++++----- ...test_mo_solve_nonhydro_stencil_11_upper.py | 10 +++++---- .../test_mo_solve_nonhydro_stencil_36.py | 3 +++ .../test_mo_solve_nonhydro_stencil_38.py | 5 +++-- .../test_mo_solve_nonhydro_stencil_39.py | 1 + .../test_mo_solve_nonhydro_stencil_40.py | 5 +++-- .../test_mo_velocity_advection_stencil_02.py | 2 ++ .../test_mo_velocity_advection_stencil_03.py | 1 + .../test_mo_velocity_advection_stencil_06.py | 5 +++-- .../test_mo_velocity_advection_stencil_10.py | 1 + 29 files changed, 86 insertions(+), 41 deletions(-) diff --git a/.gitpod.Dockerfile b/.gitpod.Dockerfile index 967ae36f2e..f79280ab05 100644 --- a/.gitpod.Dockerfile +++ b/.gitpod.Dockerfile @@ -1,7 +1,7 @@ FROM gitpod/workspace-python USER root RUN apt-get update \ - && apt-get install -y libboost-dev \ + && apt-get install -y libboost-dev mpich libnuma-dev \ && apt-get clean && rm -rf /var/cache/apt/* && rm -rf /var/lib/apt/lists/* && rm -rf /tmp/* USER gitpod RUN pyenv install 3.10.2 diff --git a/jenkins/spack-PR b/jenkins/spack-PR index e9be0ac87c..6a4c7ef9bf 100644 --- a/jenkins/spack-PR +++ b/jenkins/spack-PR @@ -13,6 +13,11 @@ triggerPhrase.each { (spack_fork, spack_branch) = parseTriggerPhrase(it) } } + +def map_srun = [ + daint: 'srun -t 03:00:00 -C gpu -A g110 -c 12 -n 1' + ] + pipeline { agent none stages { @@ -27,6 +32,7 @@ pipeline { } post { always { + archiveArtifacts artifacts: 'local_copy/spack*.txt', allowEmptyArchive: true echo 'Cleaning up workspace' deleteDir() } @@ -42,12 +48,15 @@ pipeline { } stage('Install with gt4py@main') { steps { - sh """ - . ./spack-c2sm/setup-env.sh - cd local_copy - spack env activate spack/gt4py-main - spack install -v --test=root - """ + script{ + def srun = map_srun["${NODENAME}"] + sh """ + . ./spack-c2sm/setup-env.sh + cd local_copy + spack env activate spack/gt4py-main + ${srun} spack install -v --test=root + """ + } } } } diff --git a/jenkins/spack-PR-stable b/jenkins/spack-PR-stable index b2a317a4ff..f1c72b13e0 100644 --- a/jenkins/spack-PR-stable +++ b/jenkins/spack-PR-stable @@ -13,6 +13,11 @@ triggerPhrase.each { (spack_fork, spack_branch) = parseTriggerPhrase(it) } } + +def map_srun = [ + daint: 'srun -t 03:00:00 -C gpu -A g110 -c 12 -n 1' + ] + pipeline { agent none stages { @@ -27,6 +32,7 @@ pipeline { } post { always { + archiveArtifacts artifacts: 'local_copy/spack*.txt', allowEmptyArchive: true echo 'Cleaning up workspace' deleteDir() } @@ -42,12 +48,15 @@ pipeline { } stage('Install with gt4py@stable') { steps { - sh """ - . ./spack-c2sm/setup-env.sh - cd local_copy - spack env activate spack/gt4py-stable - spack install -v --test=root - """ + script{ + def srun = map_srun["${NODENAME}"] + sh """ + . ./spack-c2sm/setup-env.sh + cd local_copy + spack env activate spack/gt4py-stable + ${srun} spack install -v --test=root + """ + } } } } diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py index 8c46ecb93e..35fbf37233 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_05.py @@ -33,4 +33,4 @@ def mo_solve_nonhydro_stencil_05( z_exner_ex_pr: Field[[CellDim, KDim], float], z_exner_ic: Field[[CellDim, KDim], float], ): - _mo_solve_nonhydro_stencil_05(wgtfac_c, z_exner_ex_pr, out=z_exner_ic) + _mo_solve_nonhydro_stencil_05(wgtfac_c, z_exner_ex_pr, out=z_exner_ic[:, 1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py index 5cb85b1374..350af63e9f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py @@ -53,5 +53,5 @@ def mo_solve_nonhydro_stencil_08( rho_ref_mc, theta_v, theta_ref_mc, - out=(rho_ic, z_rth_pr_1, z_rth_pr_2), + out=(rho_ic[:, 1:], z_rth_pr_1[:, 1:], z_rth_pr_2[:, 1:]), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py index 8ccdf40670..1dbb0f1092 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py @@ -62,5 +62,5 @@ def mo_solve_nonhydro_stencil_09( exner_pr, d_exner_dz_ref_ic, ddqz_z_half, - out=(z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c), + out=(z_theta_v_pr_ic[:, 1:], theta_v_ic[:, 1:], z_th_ddz_exner_c[:, 1:]), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index b709120573..65e7c742a7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -109,5 +109,5 @@ def mo_solve_nonhydro_stencil_10( dtime, wgt_nnow_rth, wgt_nnew_rth, - out=(rho_ic, z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c), + out=(rho_ic[:, 1:], z_theta_v_pr_ic[:, 1:], theta_v_ic[:, 1:], z_th_ddz_exner_c[:, 1:]), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py index bfb20255c9..126e59e389 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_11_upper.py @@ -47,5 +47,5 @@ def mo_solve_nonhydro_stencil_11_upper( z_rth_pr, theta_ref_ic, z_theta_v_pr_ic, - out=(z_theta_v_pr_ic, theta_v_ic), + out=(z_theta_v_pr_ic[:, -1:], theta_v_ic[:, -1:]), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py index 89299c0309..05f4293eda 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py @@ -44,4 +44,6 @@ def mo_solve_nonhydro_stencil_36( z_vt_ie: Field[[EdgeDim, KDim], float], z_kin_hor_e: Field[[EdgeDim, KDim], float], ): - _mo_solve_nonhydro_stencil_36(wgtfac_e, vn, vt, out=(vn_ie, z_vt_ie, z_kin_hor_e)) + _mo_solve_nonhydro_stencil_36( + wgtfac_e, vn, vt, out=(vn_ie[:, 1:], z_vt_ie[:, 1:], z_kin_hor_e[:, 1:]) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py index 7d49922aaf..bab81b8a1a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py @@ -37,4 +37,4 @@ def mo_solve_nonhydro_stencil_38( wgtfacq_e: Field[[EdgeDim, KDim], float], vn_ie: Field[[EdgeDim, KDim], float], ): - _mo_solve_nonhydro_stencil_38(vn, wgtfacq_e, out=vn_ie) + _mo_solve_nonhydro_stencil_38(vn, wgtfacq_e, out=vn_ie[:, -1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index 4b192e07b9..3550d2c62d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -38,4 +38,4 @@ def mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ): - _mo_solve_nonhydro_stencil_39(e_bln_c_s, z_w_concorr_me, wgtfac_c, out=w_concorr_c) + _mo_solve_nonhydro_stencil_39(e_bln_c_s, z_w_concorr_me, wgtfac_c, out=w_concorr_c[:, 1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 5b57d8f19a..dec20a7fb8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -46,4 +46,4 @@ def mo_solve_nonhydro_stencil_40( wgtfacq_c: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ): - _mo_solve_nonhydro_stencil_40(e_bln_c_s, z_w_concorr_me, wgtfacq_c, out=w_concorr_c) + _mo_solve_nonhydro_stencil_40(e_bln_c_s, z_w_concorr_me, wgtfacq_c, out=w_concorr_c[:, -1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py index 15a16dc312..8d4c673c03 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py @@ -38,4 +38,4 @@ def mo_velocity_advection_stencil_02( vn_ie: Field[[EdgeDim, KDim], float], z_kin_hor_e: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_02(wgtfac_e, vn, vt, out=(vn_ie, z_kin_hor_e)) + _mo_velocity_advection_stencil_02(wgtfac_e, vn, vt, out=(vn_ie[:, 1:], z_kin_hor_e[:, 1:])) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py index 12a7528853..1683a90d79 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_03.py @@ -34,4 +34,4 @@ def mo_velocity_advection_stencil_03( vt: Field[[EdgeDim, KDim], float], z_vt_ie: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_03(wgtfac_e, vt, out=z_vt_ie) + _mo_velocity_advection_stencil_03(wgtfac_e, vt, out=z_vt_ie[:, 1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py index 4b158fbd90..f62c74773a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py @@ -38,4 +38,4 @@ def mo_velocity_advection_stencil_06( vn: Field[[EdgeDim, KDim], float], vn_ie: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_06(wgtfacq_e, vn, out=vn_ie) + _mo_velocity_advection_stencil_06(wgtfacq_e, vn, out=vn_ie[:, -1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 4dd324603b..ab4d7f5a70 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -34,4 +34,4 @@ def mo_velocity_advection_stencil_10( wgtfac_c: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_10(z_w_concorr_mc, wgtfac_c, out=w_concorr_c) + _mo_velocity_advection_stencil_10(z_w_concorr_mc, wgtfac_c, out=w_concorr_c[:, 1:]) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py index 6d54d033aa..7007149a77 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py @@ -29,6 +29,7 @@ class TestMoSolveNonhydroStencil05(StencilTest): def reference(mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs) -> np.array: z_exner_ex_pr_offset_1 = np.roll(z_exner_ex_pr, shift=1, axis=1) z_exner_ic = wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 + z_exner_ic[:, 0] = 0 return dict(z_exner_ic=z_exner_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py index 4f9943e5de..a213a66cf6 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_08.py @@ -59,6 +59,10 @@ def reference( ) -> tuple[np.array, np.array, np.array]: rho_offset_1 = np.roll(rho, shift=1, axis=1) rho_ic = wgtfac_c * rho + (1.0 - wgtfac_c) * rho_offset_1 + rho_ic[:, 0] = 0 z_rth_pr_1 = rho - rho_ref_mc + z_rth_pr_1[:, 0] = 0 z_rth_pr_2 = theta_v - theta_ref_mc + z_rth_pr_2[:, 0] = 0 + return dict(rho_ic=rho_ic, z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py index 3939ff0d75..0695cf78c9 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_09.py @@ -43,11 +43,14 @@ def reference( vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) z_theta_v_pr_ic = wgtfac_c * z_rth_pr_2 + (1.0 - wgtfac_c) * z_rth_pr_2_offset + z_theta_v_pr_ic[:, 0] = 0 theta_v_ic = wgtfac_c * theta_v + (1 - wgtfac_c) * theta_v_offset + theta_v_ic[:, 0] = 0 z_th_ddz_exner_c = ( vwind_expl_wgt * theta_v_ic * (exner_pr_offset - exner_pr) / ddqz_z_half + z_theta_v_pr_ic * d_exner_dz_ref_ic ) + z_th_ddz_exner_c[:, 0] = 0 return dict( z_theta_v_pr_ic=z_theta_v_pr_ic, theta_v_ic=theta_v_ic, diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py index d66d399fa2..4aed2f1c5d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py @@ -18,7 +18,7 @@ mo_solve_nonhydro_stencil_10, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestMoSolveNonhydroStencil10(StencilTest): @@ -63,18 +63,22 @@ def reference( + (1 - wgtfac_c) * z_rho_tavg_m1 + z_w_backtraj * (z_rho_tavg_m1 - z_rho_tavg) ) + rho_ic[:, 0] = 0 z_theta_v_pr_mc_m1 = z_theta_tavg_m1 - theta_ref_mc_offset z_theta_v_pr_mc = z_theta_tavg - theta_ref_mc z_theta_v_pr_ic = wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 + z_theta_v_pr_ic[:, 0] = 0 theta_v_ic = ( wgtfac_c * z_theta_tavg + (1 - wgtfac_c) * z_theta_tavg_m1 + z_w_backtraj * (z_theta_tavg_m1 - z_theta_tavg) ) + theta_v_ic[:, 0] = 0 z_th_ddz_exner_c = ( vwind_expl_wgt * theta_v_ic * (exner_pr_offset - exner_pr) / ddqz_z_half + z_theta_v_pr_ic * d_exner_dz_ref_ic ) + z_th_ddz_exner_c[:, 0] = 0 return dict( rho_ic=rho_ic, z_theta_v_pr_ic=z_theta_v_pr_ic, @@ -99,10 +103,10 @@ def input_data(self, mesh): vwind_expl_wgt = random_field(mesh, CellDim) exner_pr = random_field(mesh, CellDim, KDim) d_exner_dz_ref_ic = random_field(mesh, CellDim, KDim) - rho_ic = random_field(mesh, CellDim, KDim) - z_theta_v_pr_ic = random_field(mesh, CellDim, KDim) - theta_v_ic = random_field(mesh, CellDim, KDim) - z_th_ddz_exner_c = random_field(mesh, CellDim, KDim) + rho_ic = zero_field(mesh, CellDim, KDim) + z_theta_v_pr_ic = zero_field(mesh, CellDim, KDim) + theta_v_ic = zero_field(mesh, CellDim, KDim) + z_th_ddz_exner_c = zero_field(mesh, CellDim, KDim) return dict( w=w, diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py index 891bbb585e..e50e9eade6 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_11_upper.py @@ -34,13 +34,15 @@ def reference( z_theta_v_pr_ic: np.array, **kwargs, ) -> tuple[np.array, np.array]: - z_theta_v_pr_ic = ( + z_theta_v_pr_ic_ref = np.copy(z_theta_v_pr_ic) + z_theta_v_pr_ic_ref[:, -1] = ( np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_rth_pr, shift=1, axis=1) + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_rth_pr, shift=2, axis=1) + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_rth_pr, shift=3, axis=1) - ) - theta_v_ic = theta_ref_ic + z_theta_v_pr_ic - return dict(z_theta_v_pr_ic=z_theta_v_pr_ic, theta_v_ic=theta_v_ic) + )[:, -1] + theta_v_ic = np.zeros_like(z_theta_v_pr_ic) + theta_v_ic[:, -1] = (theta_ref_ic + z_theta_v_pr_ic_ref)[:, -1] + return dict(z_theta_v_pr_ic=z_theta_v_pr_ic_ref, theta_v_ic=theta_v_ic) @pytest.fixture def input_data(self, mesh): diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py index 07cf85341b..9ab9c8568d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_36.py @@ -33,8 +33,11 @@ def reference( vt_offset_1 = np.roll(vt, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_offset_1 + vn_ie[:, 0] = 0 z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_offset_1 + z_vt_ie[:, 0] = 0 z_kin_hor_e = 0.5 * (vn**2 + vt**2) + z_kin_hor_e[:, 0] = 0 return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py index de6a9962cd..181ab847b0 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_38.py @@ -27,11 +27,12 @@ class TestMoSolveNonhydroStencil38(StencilTest): @staticmethod def reference(mesh, vn: np.array, wgtfacq_e: np.array, **kwargs) -> np.array: - vn_ie = ( + vn_ie = np.zeros_like(vn) + vn_ie[:, -1] = ( np.roll(wgtfacq_e, shift=1, axis=1) * np.roll(vn, shift=1, axis=1) + np.roll(wgtfacq_e, shift=2, axis=1) * np.roll(vn, shift=2, axis=1) + np.roll(wgtfacq_e, shift=3, axis=1) * np.roll(vn, shift=3, axis=1) - ) + )[:, -1] return dict(vn_ie=vn_ie) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index 0530136399..63e3dba809 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -38,6 +38,7 @@ def reference( z_w_concorr_mc_m0 = np.sum(e_bln_c_s * z_w_concorr_me[mesh.c2e], axis=1) z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 + w_concorr_c[:, 0] = 0 return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index 6774dc9461..dd7bc733be 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -40,11 +40,12 @@ def reference( z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) z_w_concorr_mc_m2 = np.sum(e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1) z_w_concorr_mc_m3 = np.sum(e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1) - w_concorr_c = ( + w_concorr_c = np.zeros_like(wgtfacq_c) + w_concorr_c[:, -1] = ( np.roll(wgtfacq_c, shift=1, axis=1) * z_w_concorr_mc_m1 + np.roll(wgtfacq_c, shift=2, axis=1) * z_w_concorr_mc_m2 + np.roll(wgtfacq_c, shift=3, axis=1) * z_w_concorr_mc_m3 - ) + )[:, -1] return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index 142b7e56cb..9431248f03 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -29,11 +29,13 @@ class TestMoVelocityAdvectionStencil02VnIe(StencilTest): def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 + vn_ie[:, 0] = 0 return vn_ie @staticmethod def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) + z_kin_hor_e[:, 0] = 0 return z_kin_hor_e @classmethod diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py index 6c91bd049e..cad48459fc 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py @@ -29,6 +29,7 @@ class TestMoVelocityAdvectionStencil03(StencilTest): def reference(mesh, wgtfac_e: np.array, vt: np.array, **kwargs) -> np.array: vt_k_minus_1 = np.roll(vt, shift=1, axis=1) z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 + z_vt_ie[:, 0] = 0 return dict(z_vt_ie=z_vt_ie) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py index c09d4a2040..4e68e1e05c 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py @@ -33,11 +33,12 @@ def reference(mesh, wgtfacq_e: np.array, vn: np.array, **kwargs) -> np.array: wgtfacq_e_k_minus_1 = np.roll(wgtfacq_e, shift=1, axis=1) wgtfacq_e_k_minus_2 = np.roll(wgtfacq_e, shift=2, axis=1) wgtfacq_e_k_minus_3 = np.roll(wgtfacq_e, shift=3, axis=1) - vn_ie = ( + vn_ie = np.zeros_like(vn) + vn_ie[:, -1] = ( wgtfacq_e_k_minus_1 * vn_k_minus_1 + wgtfacq_e_k_minus_2 * vn_k_minus_2 + wgtfacq_e_k_minus_3 * vn_k_minus_3 - ) + )[:, -1] return dict(vn_ie=vn_ie) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index d7d7e3d2c7..c132739d7b 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -29,6 +29,7 @@ class TestMoVelocityAdvectionStencil10(StencilTest): def reference(mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + w_concorr_c[:, 0] = 0 return dict(w_concorr_c=w_concorr_c) @pytest.fixture From 6358ce9cab6c7ae9ec476eb9dc2e4e483102ce64 Mon Sep 17 00:00:00 2001 From: Samuel Date: Thu, 21 Sep 2023 13:59:40 +0200 Subject: [PATCH 099/105] Fix typing in icon4pytools (#267) Fix static typing annotations in icon4pytools. --- base-requirements-dev.txt | 1 + .../model/atmosphere/diffusion/py.typed | 0 .../src/icon4py/model/common/grid/vertical.py | 3 +- tools/.flake8 | 4 +- tools/src/icon4pytools/f2ser/deserialise.py | 72 ++++++++----------- tools/src/icon4pytools/f2ser/parse.py | 31 ++++---- tools/src/icon4pytools/icon4pygen/backend.py | 7 +- .../icon4pygen/bindings/codegen/cpp.py | 4 +- .../icon4pygen/bindings/codegen/f90.py | 14 ++-- .../icon4pygen/bindings/codegen/header.py | 4 +- .../bindings/codegen/render/field.py | 5 +- .../icon4pygen/bindings/codegen/types.py | 4 +- .../icon4pygen/bindings/entities.py | 13 ++-- .../icon4pytools/icon4pygen/bindings/utils.py | 2 +- tools/src/icon4pytools/icon4pygen/metadata.py | 47 ++++++------ tools/src/icon4pytools/liskov/cli.py | 12 +++- .../liskov/codegen/integration/deserialise.py | 22 +++--- .../liskov/codegen/integration/generate.py | 36 ++++++---- .../liskov/codegen/integration/template.py | 39 +++++----- .../codegen/serialisation/deserialise.py | 4 +- .../liskov/codegen/serialisation/template.py | 45 ++++++------ .../liskov/codegen/shared/deserialise.py | 8 ++- .../icon4pytools/liskov/external/metadata.py | 8 ++- .../src/icon4pytools/liskov/parsing/parse.py | 7 +- .../icon4pytools/liskov/parsing/transform.py | 8 +-- .../src/icon4pytools/liskov/parsing/types.py | 3 + .../liskov/pipeline/collection.py | 7 +- tools/src/icon4pytools/py2f/cffi_utils.py | 6 +- tools/src/icon4pytools/py2f/codegen.py | 4 +- tools/src/icon4pytools/py2f/typing_utils.py | 14 ++-- .../py2f/wrappers/diffusion_wrapper.py | 7 +- .../no_deps_subroutine_example.f90 | 2 + tools/tests/f2ser/test_f2ser_codegen.py | 8 +-- 33 files changed, 236 insertions(+), 215 deletions(-) create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/py.typed diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 991b3d4732..0440e449e3 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -28,3 +28,4 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget>=3.2 +types-cffi>=1.15 diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/py.typed b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 061e6a786a..19682d7927 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -15,6 +15,7 @@ from typing import Final import numpy as np +from gt4py.next import common from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import KDim @@ -34,7 +35,7 @@ class VerticalModelParams: rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml """ - vct_a: Field[[KDim], float] + vct_a: common.Field rayleigh_damping_height: Final[float] index_of_damping_layer: Final[int32] = field(init=False) diff --git a/tools/.flake8 b/tools/.flake8 index 31cecff5ab..9bbc5e1a7f 100644 --- a/tools/.flake8 +++ b/tools/.flake8 @@ -15,7 +15,9 @@ extend-ignore = # Line too long (using Bugbear's B950 warning) E501, # Line break occurred before a binary operator - W503 + W503, + # Calling setattr with a constant attribute value + B010 exclude = .eggs, diff --git a/tools/src/icon4pytools/f2ser/deserialise.py b/tools/src/icon4pytools/f2ser/deserialise.py index 883daddc94..9aea7d72c7 100644 --- a/tools/src/icon4pytools/f2ser/deserialise.py +++ b/tools/src/icon4pytools/f2ser/deserialise.py @@ -31,40 +31,32 @@ def __init__( self.parsed = parsed self.directory = directory self.prefix = prefix - self.data = {"Savepoint": [], "Init": ..., "Import": ...} def __call__(self) -> SerialisationCodeInterface: - """Deserialise the parsed granule and returns a serialisation interface. - - Returns: - A `SerialisationInterface` object representing the deserialised data. - """ + """Deserialise the parsed granule and returns a serialisation interface.""" self._merge_out_inout_fields() - self._make_savepoints() - self._make_init_data() - self._make_imports() - return SerialisationCodeInterface(**self.data) + savepoints = self._make_savepoints() + init_data = self._make_init_data() + import_data = self._make_imports() + return SerialisationCodeInterface(Import=import_data, Init=init_data, Savepoint=savepoints) - def _make_savepoints(self) -> None: - """Create savepoints for each subroutine and intent in the parsed granule. + def _make_savepoints(self) -> list[SavepointData]: + """Create savepoints for each subroutine and intent in the parsed granule.""" + savepoints: list[SavepointData] = [] - Returns: - None. - """ for subroutine_name, intent_dict in self.parsed.subroutines.items(): for intent, var_dict in intent_dict.items(): - self._create_savepoint(subroutine_name, intent, var_dict) + savepoints.append(self._create_savepoint(subroutine_name, intent, var_dict)) - def _create_savepoint(self, subroutine_name: str, intent: str, var_dict: dict) -> None: + return savepoints + + def _create_savepoint(self, subroutine_name: str, intent: str, var_dict: dict) -> SavepointData: """Create a savepoint for the given variables. Args: subroutine_name: The name of the subroutine. intent: The intent of the fields to be serialised. var_dict: A dictionary representing the variables to be saved. - - Returns: - None. """ field_vals = {k: v for k, v in var_dict.items() if isinstance(v, dict)} fields = [ @@ -80,14 +72,12 @@ def _create_savepoint(self, subroutine_name: str, intent: str, var_dict: dict) - for var_name, var_data in field_vals.items() ] - self.data["Savepoint"].append( - SavepointData( - subroutine=subroutine_name, - intent=intent, - startln=self._get_codegen_line(var_dict["codegen_ctx"], intent), - fields=fields, - metadata=None, - ) + return SavepointData( + subroutine=subroutine_name, + intent=intent, + startln=self._get_codegen_line(var_dict["codegen_ctx"], intent), + fields=fields, + metadata=None, ) @staticmethod @@ -123,31 +113,25 @@ def _create_association(self, var_data: dict, var_name: str) -> str: ) return var_name - def _make_init_data(self) -> None: - """Create an `InitData` object and sets it to the `Init` key in the `data` dictionary. - - Returns: - None. - """ + def _make_init_data(self) -> InitData: + """Create an `InitData` object and sets it to the `Init` key in the `data` dictionary.""" first_intent_in_subroutine = [ var_dict for intent_dict in self.parsed.subroutines.values() for intent, var_dict in intent_dict.items() if intent == "in" ][0] + startln = self._get_codegen_line(first_intent_in_subroutine["codegen_ctx"], "init") - self.data["Init"] = InitData( + + return InitData( startln=startln, directory=self.directory, prefix=self.prefix, ) - def _merge_out_inout_fields(self): - """Merge the `inout` fields into the `in` and `out` fields in the `parsed` dictionary. - - Returns: - None. - """ + def _merge_out_inout_fields(self) -> None: + """Merge the `inout` fields into the `in` and `out` fields in the `parsed` dictionary.""" for _, intent_dict in self.parsed.subroutines.items(): if "inout" in intent_dict: intent_dict["in"].update(intent_dict["inout"]) @@ -155,7 +139,7 @@ def _merge_out_inout_fields(self): del intent_dict["inout"] @staticmethod - def _get_codegen_line(ctx: CodegenContext, intent: str): + def _get_codegen_line(ctx: CodegenContext, intent: str) -> int: if intent == "in": return ctx.last_declaration_ln elif intent == "out": @@ -165,5 +149,5 @@ def _get_codegen_line(ctx: CodegenContext, intent: str): else: raise ValueError(f"Unrecognized intent: {intent}") - def _make_imports(self): - self.data["Import"] = ImportData(startln=self.parsed.last_import_ln) + def _make_imports(self) -> ImportData: + return ImportData(startln=self.parsed.last_import_ln) diff --git a/tools/src/icon4pytools/f2ser/parse.py b/tools/src/icon4pytools/f2ser/parse.py index a68fe99003..85780d91a9 100644 --- a/tools/src/icon4pytools/f2ser/parse.py +++ b/tools/src/icon4pytools/f2ser/parse.py @@ -38,7 +38,7 @@ class CodegenContext: end_subroutine_ln: int -ParsedSubroutines = dict[str, dict[str, dict[str, any] | CodegenContext]] +ParsedSubroutines = dict[str, dict[str, dict[str, Any]]] @dataclass @@ -69,11 +69,11 @@ def __call__(self) -> ParsedGranule: last_import_ln = self._find_last_fortran_use_statement() return ParsedGranule(subroutines=subroutines, last_import_ln=last_import_ln) - def _find_last_fortran_use_statement(self) -> Optional[int]: + def _find_last_fortran_use_statement(self) -> int: """Find the line number of the last Fortran USE statement in the code. Returns: - int: the line number of the last USE statement, or None if no USE statement is found. + int: the line number of the last USE statement. """ # Reverse the order of the lines so we can search from the end code = self._read_code_from_file() @@ -81,16 +81,14 @@ def _find_last_fortran_use_statement(self) -> Optional[int]: code_lines.reverse() # Look for the last USE statement - use_ln = None for i, line in enumerate(code_lines): if line.strip().lower().startswith("use"): use_ln = len(code_lines) - i if i > 0 and code_lines[i - 1].strip().lower() == "#endif": # If the USE statement is preceded by an #endif statement, return the line number after the #endif statement - return use_ln + 1 - else: - return use_ln - return None + use_ln += 1 + return use_ln + raise ParsingError("Could not find any USE statements.") def _read_code_from_file(self) -> str: """Read the content of the granule and returns it as a string.""" @@ -98,7 +96,7 @@ def _read_code_from_file(self) -> str: code = f.read() return code - def parse_subroutines(self): + def parse_subroutines(self) -> dict: subroutines = self._extract_subroutines(crack(self.granule_path)) variables_grouped_by_intent = { name: self._extract_intent_vars(routine) for name, routine in subroutines.items() @@ -263,7 +261,7 @@ def _combine_types(derived_type_vars: dict, intrinsic_type_vars: dict) -> dict: combined[subroutine_name][intent].update(new_vars) return combined - def _update_with_codegen_lines(self, parsed_types: dict) -> dict: + def _update_with_codegen_lines(self, parsed_types: dict[str, Any]) -> dict[str, Any]: """Update the parsed_types dictionary with the line numbers for codegen. Args: @@ -285,9 +283,6 @@ def _get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: Args: subroutine_name (str): Name of the subroutine to look for in the code. - - Returns: - CodegenContext: Object containing the line number of the last declaration statement and the line number of the last line of the code before the end of the given subroutine. """ code = self._read_code_from_file() @@ -312,7 +307,7 @@ def _get_subroutine_lines(self, subroutine_name: str) -> CodegenContext: return CodegenContext(first_declaration_ln, last_declaration_ln, pre_end_subroutine_ln) @staticmethod - def _find_subroutine_lines(code: str, subroutine_name: str) -> tuple[int]: + def _find_subroutine_lines(code: str, subroutine_name: str) -> tuple[int, int]: """Find line numbers of a subroutine within a code block. Args: @@ -327,7 +322,7 @@ def _find_subroutine_lines(code: str, subroutine_name: str) -> tuple[int]: start_match = re.search(start_subroutine_pattern, code) end_match = re.search(end_subroutine_pattern, code) if start_match is None or end_match is None: - return None + raise ParsingError(f"Could not find {start_match} or {end_match}") start_subroutine_ln = code[: start_match.start()].count("\n") + 1 end_subroutine_ln = code[: end_match.start()].count("\n") + 1 return start_subroutine_ln, end_subroutine_ln @@ -335,7 +330,7 @@ def _find_subroutine_lines(code: str, subroutine_name: str) -> tuple[int]: @staticmethod def _find_variable_declarations( code: str, start_subroutine_ln: int, end_subroutine_ln: int - ) -> list: + ) -> list[int]: """Find line numbers of variable declarations within a code block. Args: @@ -371,8 +366,8 @@ def _find_variable_declarations( @staticmethod def _get_variable_declaration_bounds( - declaration_pattern_lines: list, start_subroutine_ln: int - ) -> tuple: + declaration_pattern_lines: list[int], start_subroutine_ln: int + ) -> tuple[int, int]: """Return the line numbers of the bounds for a variable declaration block. Args: diff --git a/tools/src/icon4pytools/icon4pygen/backend.py b/tools/src/icon4pytools/icon4pygen/backend.py index 01b917afd1..6522a41bd8 100644 --- a/tools/src/icon4pytools/icon4pygen/backend.py +++ b/tools/src/icon4pytools/icon4pygen/backend.py @@ -95,7 +95,6 @@ def _is_size_param(param: itir.Sym) -> bool: @staticmethod def _missing_domain_params(params: List[itir.Sym]) -> Iterable[itir.Sym]: """Get domain limit params that are not present in param list.""" - return map( - lambda p: itir.Sym(id=p), - filter(lambda s: s not in map(lambda p: p.id, params), _DOMAIN_ARGS), - ) + param_ids = [p.id for p in params] + missing_args = [s for s in _DOMAIN_ARGS if s not in param_ids] + return (itir.Sym(id=p) for p in missing_args) diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py index ba6951f572..7e1369a929 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/cpp.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path -from typing import Sequence +from typing import Any, Sequence from gt4py import eve from gt4py.eve.codegen import JinjaTemplate as as_jinja @@ -678,7 +678,7 @@ def _get_field_data(self) -> tuple: ) return fields, offsets - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: fields, offsets = self._get_field_data() offset_renderer = GpuTriMeshOffsetRenderer(self.offsets) diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/f90.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/f90.py index a3cc0a3e41..5bc1418e71 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/f90.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/f90.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path -from typing import Sequence, Union +from typing import Any, Sequence, Union from gt4py import eve from gt4py.eve import Node @@ -214,7 +214,7 @@ class F90RunFun(eve.Node): params: F90EntityList = eve.datamodels.field(init=False) binds: F90EntityList = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: param_fields = [F90Field(name=field.name) for field in self.all_fields] + [ F90Field(name=name) for name in _DOMAIN_ARGS ] @@ -242,7 +242,7 @@ class F90RunAndVerifyFun(eve.Node): params: F90EntityList = eve.datamodels.field(init=False) binds: F90EntityList = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: param_fields = ( [F90Field(name=field.name) for field in self.all_fields] + [F90Field(name=field.name, suffix="before") for field in self.out_fields] @@ -295,7 +295,7 @@ class F90SetupFun(Node): params: F90EntityList = eve.datamodels.field(init=False) binds: F90EntityList = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: param_fields = [ F90Field(name=name) for name in [ @@ -346,7 +346,7 @@ class F90WrapRunFun(Node): run_ver_params: F90EntityList = eve.datamodels.field(init=False) run_params: F90EntityList = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: param_fields = ( [F90Field(name=field.name) for field in self.all_fields] + [F90Field(name=field.name, suffix="before") for field in self.out_fields] @@ -457,7 +457,7 @@ class F90WrapSetupFun(Node): vert_conditionals: F90EntityList = eve.datamodels.field(init=False) setup_params: F90EntityList = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: param_fields = [ F90Field(name=name) for name in [ @@ -534,7 +534,7 @@ class F90File(Node): wrap_run_fun: F90WrapRunFun = eve.datamodels.field(init=False) wrap_setup_fun: F90WrapSetupFun = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: all_fields = self.fields out_fields = [field for field in self.fields if field.intent.out] tol_fields = [field for field in out_fields if not field.is_integral()] diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/header.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/header.py index 045cea7827..35775b5a39 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/header.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/header.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path -from typing import Sequence +from typing import Any, Sequence from gt4py import eve from gt4py.eve import Node @@ -149,7 +149,7 @@ class CppHeaderFile(Node): setupFunc: CppSetupFuncDeclaration = eve.datamodels.field(init=False) freeFunc: CppFreeFunc = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: output_fields = [field for field in self.fields if field.intent.out] tolerance_fields = [field for field in output_fields if not field.is_integral()] diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/render/field.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/render/field.py index e9106fba0e..79c67cd819 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/render/field.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/render/field.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from dataclasses import dataclass +from typing import cast from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, @@ -19,6 +20,7 @@ ) from icon4pytools.icon4pygen.bindings.codegen.types import FieldEntity from icon4pytools.icon4pygen.bindings.exceptions import BindingsRenderingException +from icon4pytools.icon4pygen.bindings.locations import ChainedLocation, CompoundLocation @dataclass(frozen=True) @@ -85,7 +87,8 @@ def render_stride_type(self) -> str: if self.entity.is_dense(): return _strides[str(self.entity.location)] elif self.entity.is_sparse() or self.entity.is_compound(): - return _strides[str(self.entity.location[0])] # type: ignore + location = cast(ChainedLocation | CompoundLocation, self.entity.location) + return _strides[str(location[0])] else: raise BindingsRenderingException("stride type called on scalar") diff --git a/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py b/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py index 87ffad337c..e7b63c2230 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/codegen/types.py @@ -13,7 +13,7 @@ import abc from abc import ABC from dataclasses import dataclass -from typing import Union +from typing import Optional, Union from gt4py.next.type_system import type_specifications as ts @@ -36,7 +36,7 @@ class FieldEntity(ABC): intent: FieldIntent has_vertical_dimension: bool includes_center: bool - location: ChainedLocation | CompoundLocation | BasicLocation | None + location: Optional[ChainedLocation | CompoundLocation | BasicLocation] @abc.abstractmethod def is_sparse(self) -> bool: diff --git a/tools/src/icon4pytools/icon4pygen/bindings/entities.py b/tools/src/icon4pytools/icon4pygen/bindings/entities.py index cad08c7c0d..3d74f4be4d 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/entities.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/entities.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Union +from typing import cast from gt4py.eve import Node from gt4py.next.ffront import program_ast as past @@ -63,7 +63,7 @@ def _includes_center(chain: str) -> bool: return False @staticmethod - def _handle_source(chain: str) -> Union[BasicLocation, CompoundLocation]: + def _handle_source(chain: str) -> BasicLocation | CompoundLocation: if chain.endswith("O"): chain = chain[:-1] @@ -78,7 +78,7 @@ def _handle_source(chain: str) -> Union[BasicLocation, CompoundLocation]: @staticmethod def _make_target( - chain: str, source: Union[BasicLocation, CompoundLocation] + chain: str, source: BasicLocation | CompoundLocation ) -> tuple[BasicLocation, ChainedLocation]: if chain.endswith("O"): chain = chain[:-1] @@ -130,17 +130,18 @@ def rank(self) -> int: return rank def get_num_neighbors(self) -> int: - if not (self.is_sparse() or self.is_compound()): + if not (self.is_sparse() or self.is_compound()) or self.location is None: raise BindingsTypeConsistencyException( "num nbh only defined for sparse or compound fields" ) - return calc_num_neighbors(self.location.to_dim_list(), self.includes_center) # type: ignore + location = cast(ChainedLocation | CompoundLocation, self.location) + return calc_num_neighbors(location.to_dim_list(), self.includes_center) @staticmethod def _extract_field_type(field: past.DataSymbol) -> ts.ScalarKind: """Handle extraction of field types for different fields e.g. Scalar.""" if not isinstance(field.type, ts.FieldType): - return field.type.kind # type: ignore + return field.type.kind # type: ignore[union-attr] return field.type.dtype.kind @staticmethod diff --git a/tools/src/icon4pytools/icon4pygen/bindings/utils.py b/tools/src/icon4pytools/icon4pygen/bindings/utils.py index 5d5cb72c1e..081a560c28 100644 --- a/tools/src/icon4pytools/icon4pygen/bindings/utils.py +++ b/tools/src/icon4pytools/icon4pygen/bindings/utils.py @@ -48,7 +48,7 @@ def format_fortran_code(source: str) -> str: if fprettify_path is None: bin_path = Path(PYTHON_PATH).parent - fprettify_path = bin_path / "fprettify" + fprettify_path = str(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() diff --git a/tools/src/icon4pytools/icon4pygen/metadata.py b/tools/src/icon4pytools/icon4pygen/metadata.py index 773c3d2b27..fd9869016e 100644 --- a/tools/src/icon4pytools/icon4pygen/metadata.py +++ b/tools/src/icon4pytools/icon4pygen/metadata.py @@ -13,7 +13,6 @@ from __future__ import annotations import importlib -import inspect import types from dataclasses import dataclass from typing import Any, Optional, TypeGuard @@ -69,8 +68,21 @@ def is_list_of_names(obj: Any) -> TypeGuard[list[past.Name]]: return isinstance(obj, list) and all(isinstance(i, past.Name) for i in obj) -def _ignore_subscript(node: past.Name | past.Subscript) -> past.Name: - return node if isinstance(node, past.Name) else node.value +def is_name(node: past.Expr) -> TypeGuard[past.Name]: + return isinstance(node, past.Name) + + +def is_subscript(node: past.Expr) -> TypeGuard[past.Subscript]: + return isinstance(node, past.Subscript) + + +def _ignore_subscript(node: past.Expr) -> past.Name: + if is_name(node): + return node + elif is_subscript(node): + return node.value + else: + raise Exception("Need only past.Name in output kwargs.") def _get_field_infos(fvprog: Program) -> dict[str, FieldInfo]: @@ -87,7 +99,7 @@ def _get_field_infos(fvprog: Program) -> dict[str, FieldInfo]: else [_ignore_subscript(out_arg)] ) assert all(isinstance(f, past.Name) for f in output_fields) - output_arg_ids = set(arg.id for arg in output_fields) # type: ignore + output_arg_ids = set(arg.id for arg in output_fields) domain_arg_ids = _get_domain_arg_ids(fvprog) @@ -104,7 +116,7 @@ def _get_field_infos(fvprog: Program) -> dict[str, FieldInfo]: return fields -def _get_domain_arg_ids(fvprog: Program) -> list | set[str]: +def _get_domain_arg_ids(fvprog: Program) -> set[Optional[eve.concepts.SymbolRef]]: """Collect all argument names that are used within the 'domain' keyword argument.""" domain_arg_ids = [] if "domain" in fvprog.past_node.body[0].kwargs.keys(): @@ -114,8 +126,7 @@ def _get_domain_arg_ids(fvprog: Program) -> list | set[str]: for arg_elt in arg.elts: if isinstance(arg_elt, past.Name): domain_arg_ids.append(arg_elt.id) - domain_arg_ids = set(domain_arg_ids) - return domain_arg_ids + return set(domain_arg_ids) def import_definition(name: str) -> Program | FieldOperator | types.FunctionType: @@ -215,23 +226,15 @@ def scan_for_offsets(fvprog: Program) -> list[eve.concepts.SymbolRef]: def get_stencil_info( - fencil_def: Program | FieldOperator | types.FunctionType, is_global: bool = False + fencil_def: Program | FieldOperator | types.FunctionType | FendefDispatcher, + is_global: bool = False, ) -> StencilInfo: """Generate StencilInfo dataclass from a fencil definition.""" - if isinstance(fencil_def, FendefDispatcher): - # this branch can be removed once we don't have plain itir programs - params = inspect.signature(fencil_def.function).parameters - num_params = len(params) - itir = fencil_def.itir(*[None] * num_params) # TODO do this in GT4Py? - offsets = fencil_def.offsets - fields = fencil_def.metadata - column_axis = None - else: - fvprog = get_fvprog(fencil_def) - offsets = scan_for_offsets(fvprog) - itir = fvprog.itir - fields = _get_field_infos(fvprog) - column_axis = fvprog._column_axis + fvprog = get_fvprog(fencil_def) + offsets = scan_for_offsets(fvprog) + itir = fvprog.itir + fields = _get_field_infos(fvprog) + column_axis = fvprog._column_axis offset_provider = {} for offset in offsets: diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index e6685c3ad5..dd6eb7876b 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -29,7 +29,7 @@ @click.group(invoke_without_command=True) @click.pass_context -def main(ctx): +def main(ctx: click.Context) -> None: """Command line interface for interacting with the ICON-Liskov DSL Preprocessor.""" if ctx.invoked_subcommand is None: raise MissingCommandError( @@ -64,7 +64,13 @@ def main(ctx): "output_path", type=click.Path(dir_okay=False, resolve_path=True, path_type=pathlib.Path), ) -def integrate(input_path, output_path, fused, profile, metadatagen): +def integrate( + input_path: pathlib.Path, + output_path: pathlib.Path, + fused: bool, + profile: bool, + metadatagen: bool, +) -> None: mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) iface_gt4py = process_stencils(iface, fused) @@ -94,7 +100,7 @@ def integrate(input_path, output_path, fused, profile, metadatagen): "output_path", type=click.Path(dir_okay=False, resolve_path=True, path_type=pathlib.Path), ) -def serialise(input_path, output_path, multinode): +def serialise(input_path: pathlib.Path, output_path: pathlib.Path, multinode: bool) -> None: mode = "serialisation" iface = parse_fortran_file(input_path, output_path, mode) run_code_generation(input_path, output_path, mode, iface, multinode=multinode) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index f4118c9e95..fd3c9aecb5 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Any, Optional, Protocol, Type +from typing import Any, Optional, Protocol, Sequence, Type, cast import icon4pytools.liskov.parsing.parse import icon4pytools.liskov.parsing.types as ts @@ -82,7 +82,7 @@ def _extract_boolean_kwarg( return False -def pop_item_from_dict(dictionary: dict, key: str, default_value: str) -> str: +def pop_item_from_dict(dictionary: dict, key: str, default_value: Optional[str]) -> str: return dictionary.pop(key, default_value) @@ -152,7 +152,7 @@ class StartCreateDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartCreate dtype: Type[StartCreateData] = StartCreateData - def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData]: + def __call__(self, parsed: ts.ParsedDict) -> list[StartCreateData] | Type[UnusedDirective]: deserialised = [] extracted = extract_directive(parsed["directives"], self.directive_cls) @@ -261,10 +261,10 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: class StartStencilDataFactoryBase(DataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = None - dtype: Type[StartFusedStencilData] = None + directive_cls: Type[ts.ParsedDirective] + dtype: Type[StartStencilData | StartFusedStencilData] - def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: + def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData | StartFusedStencilData]: field_dimensions = flatten_list_of_dicts( [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] ) @@ -276,8 +276,8 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: def create_stencil_data( self, parsed: ts.ParsedDict, - field_dimensions: list[dict[str, Any]], - directives: list[ts.ParsedDirective], + field_dimensions: dict[str, Any], + directives: Sequence[ts.ParsedDirective], directive_cls: Type[ts.ParsedDirective], dtype: Type[StartStencilData | StartFusedStencilData], ) -> list[StartStencilData | StartFusedStencilData]: @@ -420,10 +420,8 @@ def __call__(self, parsed: ts.ParsedDict) -> list[InsertData]: deserialised = [] extracted = extract_directive(parsed["directives"], self.directive_cls) for i, directive in enumerate(extracted): - content = parsed["content"]["Insert"][i] - deserialised.append( - self.dtype(startln=directive.startln, content=content) # type: ignore - ) + content = cast(str, parsed["content"]["Insert"][i]) + deserialised.append(self.dtype(startln=directive.startln, content=content)) return deserialised diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 40324d138f..0603e226f4 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Any +from typing import Any, Sequence, TypeGuard, Union from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.codegen.integration.interface import ( @@ -59,6 +59,10 @@ logger = setup_logger(__name__) +def _is_sequence(value: Union[Sequence[Any], UnusedDirective]) -> TypeGuard[Sequence[Any]]: + return isinstance(value, Sequence) + + class IntegrationCodeGenerator(CodeGenerator): def __init__( self, @@ -204,7 +208,9 @@ def _generate_end_fused_stencil(self) -> None: def _generate_delete(self) -> None: """Generate f90 integration code for delete section.""" - if self.interface.StartDelete != UnusedDirective: + if isinstance(self.interface.StartDelete, Sequence) and isinstance( + self.interface.EndDelete, Sequence + ): logger.info("Generating DELETE statement.") for start, end in zip( self.interface.StartDelete, self.interface.EndDelete, strict=True @@ -227,13 +233,13 @@ def _generate_imports(self) -> None: ImportsStatement, ImportsStatementGenerator, self.interface.Imports.startln, - stencils=self.interface.StartStencil + self.interface.StartFusedStencil, + stencils=[*self.interface.StartStencil, *self.interface.StartFusedStencil], ) def _generate_create(self) -> None: """Generate f90 code for OpenACC DATA CREATE statements.""" - if self.interface.StartCreate != UnusedDirective: - for startcreate in self.interface.StartCreate: # type: ignore + if _is_sequence(self.interface.StartCreate): + for startcreate in self.interface.StartCreate: logger.info("Generating DATA CREATE statement.") self._generate( StartCreateStatement, @@ -242,8 +248,8 @@ def _generate_create(self) -> None: extra_fields=startcreate.extra_fields, ) - if self.interface.EndCreate != UnusedDirective: - for endcreate in self.interface.EndCreate: # type: ignore + if _is_sequence(self.interface.EndCreate): + for endcreate in self.interface.EndCreate: self._generate( EndCreateStatement, EndCreateStatementGenerator, @@ -252,16 +258,16 @@ def _generate_create(self) -> None: def _generate_endif(self) -> None: """Generate f90 code for endif statements.""" - if self.interface.EndIf != UnusedDirective: - for endif in self.interface.EndIf: # type: ignore + if _is_sequence(self.interface.EndIf): + for endif in self.interface.EndIf: logger.info("Generating ENDIF statement.") self._generate(EndIfStatement, EndIfStatementGenerator, endif.startln) def _generate_profile(self) -> None: """Generate additional nvtx profiling statements.""" if self.profile: - if self.interface.StartProfile != UnusedDirective: - for start in self.interface.StartProfile: # type: ignore + if _is_sequence(self.interface.StartProfile): + for start in self.interface.StartProfile: logger.info("Generating nvtx start statement.") self._generate( StartProfileStatement, @@ -270,8 +276,8 @@ def _generate_profile(self) -> None: name=start.name, ) - if self.interface.EndProfile != UnusedDirective: - for end in self.interface.EndProfile: # type: ignore + if _is_sequence(self.interface.EndProfile): + for end in self.interface.EndProfile: logger.info("Generating nvtx end statement.") self._generate( EndProfileStatement, @@ -281,8 +287,8 @@ def _generate_profile(self) -> None: def _generate_insert(self) -> None: """Generate free form statement from insert directive.""" - if self.interface.Insert != UnusedDirective: - for insert in self.interface.Insert: # type: ignore + if _is_sequence(self.interface.Insert): + for insert in self.interface.Insert: logger.info("Generating free form statement.") self._generate( InsertStatement, diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 2f154f82d2..046790a9cf 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -13,7 +13,7 @@ import re from dataclasses import asdict -from typing import Optional +from typing import Any, Collection, Optional import gt4py.eve as eve from gt4py.eve.codegen import JinjaTemplate as as_jinja @@ -117,7 +117,7 @@ class EndStencilStatement(EndBasicStencilStatement): noprofile: Optional[bool] noaccenddata: Optional[bool] - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args, **kwargs) -> None: all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) self.name = self.stencil_data.name @@ -141,26 +141,29 @@ class BaseEndStencilStatementGenerator(TemplatedGenerator): """ ) - OutputFields = as_jinja( - """ - {%- for field in _this_node.fields %} - {{ field.variable }}={{ field.association }},& - {{ field.variable }}_before={{ field.variable }}_before{{ field.rh_index }},& - {%- endfor %} - """ - ) + def visit_OutputFields(self, out: OutputFields) -> str | Collection[str]: + for f in out.fields: + if f.dims is None: + continue - def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore - for f in out.fields: # type: ignore idx = render_index(f.dims) split_idx = idx.split(",") if len(split_idx) >= 3: split_idx[-1] = "1" - f.rh_index = enclose_in_parentheses(",".join(split_idx)) + setattr(f, "rh_index", enclose_in_parentheses(",".join(split_idx))) return self.generic_visit(out) + OutputFields = as_jinja( + """ + {%- for field in _this_node.fields %} + {{ field.variable }}={{ field.association }},& + {{ field.variable }}_before={{ field.variable }}_before{{ field.rh_index }},& + {%- endfor %} + """ + ) + ToleranceFields = as_jinja( """ {%- if _this_node.fields|length < 1 -%} @@ -269,7 +272,7 @@ class DeclareStatement(eve.Node): declare_data: DeclareData declarations: list[Declaration] = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: self.declarations = [ Declaration(variable=k, association=v) for k, v in self.declare_data.declarations.items() @@ -292,7 +295,7 @@ class StartStencilStatement(eve.Node): profile: bool copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" @@ -302,7 +305,7 @@ class StartFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" @@ -312,7 +315,7 @@ class EndFusedStencilStatement(EndBasicStencilStatement): stencil_data: StartFusedStencilData copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) @@ -393,7 +396,7 @@ class ImportsStatement(eve.Node): stencils: list[BaseStartStencilData] stencil_names: list[str] = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore + def __post_init__(self, *args: Any, **kwargs: Any) -> None: self.stencil_names = sorted(set([stencil.name for stencil in self.stencils])) diff --git a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py index 9db5533201..3e711ccaf5 100644 --- a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py @@ -190,8 +190,8 @@ def _get_timestep_variables(stencil_name: str) -> dict: return timestep_variables - def _find_repeated_stencils(self, content): - stencil_names = {} + def _find_repeated_stencils(self, content: dict) -> set[str]: + stencil_names: dict[str, str] = {} repeated_names = [] for stencil in content["StartStencil"]: name = stencil["name"] diff --git a/tools/src/icon4pytools/liskov/codegen/serialisation/template.py b/tools/src/icon4pytools/liskov/codegen/serialisation/template.py index 8c9f8d5325..75c9cc2f60 100644 --- a/tools/src/icon4pytools/liskov/codegen/serialisation/template.py +++ b/tools/src/icon4pytools/liskov/codegen/serialisation/template.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from dataclasses import asdict -from typing import Optional +from typing import Any, Collection, Optional import gt4py.eve as eve from gt4py.eve.codegen import JinjaTemplate as as_jinja @@ -35,11 +35,11 @@ class StandardFields(eve.Node): fields: list[Field] -class DecomposedFields(StandardFields): +class DecomposedFieldsAllocNode(StandardFields): ... -class DecomposedFieldDeclarations(DecomposedFields): +class DecomposedFieldDeclarations(DecomposedFieldsAllocNode): ... @@ -48,14 +48,14 @@ class SavepointStatement(eve.Node): init: Optional[InitData] = eve.datamodels.field(default=None) multinode: bool standard_fields: StandardFields = eve.datamodels.field(init=False) - decomposed_fields: DecomposedFields = eve.datamodels.field(init=False) + decomposed_fields: DecomposedFieldsAllocNode = eve.datamodels.field(init=False) decomposed_field_declarations: DecomposedFieldDeclarations = eve.datamodels.field(init=False) - def __post_init__(self): + def __post_init__(self, *args: Any, **kwargs: Any) -> None: self.standard_fields = StandardFields( fields=[Field(**asdict(f)) for f in self.savepoint.fields if not f.decomposed] ) - self.decomposed_fields = DecomposedFields( + self.decomposed_fields = DecomposedFieldsAllocNode( fields=[Field(**asdict(f)) for f in self.savepoint.fields if f.decomposed] ) self.decomposed_field_declarations = DecomposedFieldDeclarations( @@ -105,7 +105,25 @@ class SavepointStatementGenerator(TemplatedGenerator): """ ) - DecomposedFields = as_jinja( + def visit_DecomposedFieldsAllocNode( + self, node: DecomposedFieldsAllocNode + ) -> str | Collection[str]: + def generate_size_strings(dim_list: list[str], var_name: str) -> list[str]: + size_strings = [] + for i in range(len(dim_list)): + size_strings.append(f"size({var_name}, {i + 1})") + return size_strings + + for f in node.fields: + if f.dimension is None: + raise Exception("No dimension found in `DecomposedField` {node}") + + f.variable = f.variable.replace(f"_{f.ptr_var}", "") + setattr(f, "alloc_dims", ",".join(generate_size_strings(f.dimension, f.variable))) + + return self.generic_visit(node) + + DecomposedFieldsAllocNode = as_jinja( """ {% for f in _this_node.fields %} !$ser verbatim allocate({{ f.variable }}_{{ f.ptr_var}}({{ f.alloc_dims }})) @@ -116,19 +134,6 @@ class SavepointStatementGenerator(TemplatedGenerator): """ ) - def visit_DecomposedFields(self, node: DecomposedFields): - def generate_size_strings(colon_list, var_name): - size_strings = [] - for i in range(len(colon_list)): - size_strings.append(f"size({var_name}, {i + 1})") - return size_strings - - for f in node.fields: - f.variable = f.variable.replace(f"_{f.ptr_var}", "") - f.alloc_dims = ",".join(generate_size_strings(f.dimension, f.variable)) - - return self.generic_visit(node) - class ImportStatement(eve.Node): ... diff --git a/tools/src/icon4pytools/liskov/codegen/shared/deserialise.py b/tools/src/icon4pytools/liskov/codegen/shared/deserialise.py index e71ef9fe3f..b53c02afe2 100644 --- a/tools/src/icon4pytools/liskov/codegen/shared/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/shared/deserialise.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Callable +from typing import Callable, Type import icon4pytools.liskov.parsing.types as ts from icon4pytools.common.logger import setup_logger @@ -25,9 +25,11 @@ class Deserialiser(Step): _FACTORIES: dict[str, Callable] = {} - _INTERFACE_TYPE: SerialisationCodeInterface | IntegrationCodeInterface + _INTERFACE_TYPE: Type[SerialisationCodeInterface | IntegrationCodeInterface] - def __call__(self, directives: ts.ParsedDict): + def __call__( + self, directives: ts.ParsedDict + ) -> SerialisationCodeInterface | IntegrationCodeInterface: """Deserialises parsed directives into an Interface object. Args: diff --git a/tools/src/icon4pytools/liskov/external/metadata.py b/tools/src/icon4pytools/liskov/external/metadata.py index b9d6728fdf..988056fda0 100644 --- a/tools/src/icon4pytools/liskov/external/metadata.py +++ b/tools/src/icon4pytools/liskov/external/metadata.py @@ -30,8 +30,14 @@ def generated_on(self) -> str: def cli_params(self) -> dict[str, Any]: try: ctx = click.get_current_context() + if ctx is None: + raise MissingClickContextError("No active Click context found.") + params = ctx.params.copy() - params.update(ctx.parent.params) + + if ctx.parent is not None: + params.update(ctx.parent.params) + return params except Exception as e: raise MissingClickContextError( diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index 7d60812b08..c4052fa56c 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -72,7 +72,7 @@ def _determine_type( found = False for directive in SUPPORTED_DIRECTIVES: if directive.pattern in raw.string: - typed.append(directive(raw.string, raw.startln, raw.endln)) # type: ignore + typed.append(directive(raw.string, raw.startln, raw.endln)) found = True break if not found: @@ -83,10 +83,7 @@ def _determine_type( def _preprocess(self, directives: Sequence[ts.ParsedDirective]) -> Sequence[ts.ParsedDirective]: """Preprocess the directives by removing unnecessary characters and formatting the directive strings.""" - return [ - d.__class__(self._clean_string(d.string), d.startln, d.endln) # type: ignore - for d in directives - ] + return [d.__class__(self._clean_string(d.string), d.startln, d.endln) for d in directives] def _run_validation_passes(self, preprocessed: Sequence[ts.ParsedDirective]) -> None: """Run validation passes on the directives.""" diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 12fd9c0478..3537a2dc63 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later from typing import Any -import icon4pytools.liskov.parsing.types as ts from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.codegen.integration.interface import ( EndDeleteData, @@ -24,6 +23,7 @@ StartStencilData, UnusedDirective, ) +from icon4pytools.liskov.codegen.shared.types import CodeGenInput from icon4pytools.liskov.pipeline.definition import Step @@ -35,7 +35,7 @@ def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: self.parsed = parsed self.fused = fused - def __call__(self, data: Any = None) -> ts.ParsedDict: + def __call__(self, data: Any = None) -> IntegrationCodeInterface: """Transform stencils in the parse tree based on the 'fused' flag, transforming or removing as necessary. This method processes stencils present in the 'parsed' object according to the 'fused' @@ -46,7 +46,7 @@ def __call__(self, data: Any = None) -> ts.ParsedDict: data (Any): Optional data to be passed. Default is None. Returns: - ts.ParsedDict: The parsed directives along with any modifications applied. + IntegrationCodeInterface: The interface object along with any transformations applied. """ if self.fused: logger.info("Transforming stencils for deletion.") @@ -102,7 +102,7 @@ def _create_delete_directives( directive.append(cls(startln=param.startln)) setattr(self.parsed, attr, directive) - def _remove_stencils(self, stencils_to_remove: list[StartStencilData | EndStencilData]) -> None: + def _remove_stencils(self, stencils_to_remove: list[CodeGenInput]) -> None: attributes_to_modify = ["StartStencil", "EndStencil"] for attr_name in attributes_to_modify: diff --git a/tools/src/icon4pytools/liskov/parsing/types.py b/tools/src/icon4pytools/liskov/parsing/types.py index 0c69f5475c..8889e5de65 100644 --- a/tools/src/icon4pytools/liskov/parsing/types.py +++ b/tools/src/icon4pytools/liskov/parsing/types.py @@ -25,6 +25,9 @@ class ParsedDirective(Protocol): pattern: str regex: str + def __init__(self, string: str, startln: int, endln: int) -> None: + ... + @property def type_name(self) -> str: ... diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 8af70e54ac..4d458ad1cc 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path +from typing import Any from icon4pytools.liskov.codegen.integration.deserialise import IntegrationCodeDeserialiser from icon4pytools.liskov.codegen.integration.generate import IntegrationCodeGenerator @@ -41,7 +42,7 @@ def parse_fortran_file( input_filepath: Path, output_filepath: Path, deserialiser_type: str, - **kwargs, + **kwargs: Any, ) -> list[Step]: """Execute a pipeline to parse and deserialize directives from a file. @@ -91,8 +92,8 @@ def run_code_generation( input_filepath: Path, output_filepath: Path, codegen_type: str, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> list[Step]: """Execute a pipeline to generate and write code. diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py index c251831322..9627e37080 100644 --- a/tools/src/icon4pytools/py2f/cffi_utils.py +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -31,7 +31,7 @@ class CffiMethod: - _registry = {} + _registry: dict[str, list[str]] = {} @classmethod def register(cls, func): @@ -108,7 +108,7 @@ def to_fields(dim_sizes: dict[Dimension, int]): ffi = cffi.FFI() dim_sizes = dim_sizes - def _dim_sizes(dims: list[Dimension]) -> tuple[int, int]: + def _dim_sizes(dims: list[Dimension]) -> tuple[int | None, int | None]: """Extract the size of dimension from a dictionary.""" v_size = None h_size = None @@ -140,7 +140,7 @@ def _unpack(ptr, size_h, size_v, dtype) -> np.ndarray: # TODO (magdalena) fix dtype handling use SCALARTYPE? mem_size = ffi.sizeof(c_type) mem_size = np.dtype(c_type).itemsize - ar = np.frombuffer( + ar = np.frombuffer( # type: ignore[call-overload] ffi.buffer(ptr, length * mem_size), dtype=np.dtype(c_type), count=-1, diff --git a/tools/src/icon4pytools/py2f/codegen.py b/tools/src/icon4pytools/py2f/codegen.py index 502b204922..1025552005 100644 --- a/tools/src/icon4pytools/py2f/codegen.py +++ b/tools/src/icon4pytools/py2f/codegen.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +from pathlib import Path from typing import Sequence from gt4py.eve import Node, codegen @@ -140,6 +140,6 @@ def generate_c_header(plugin: CffiPlugin) -> str: return codegen.format_source("cpp", generated_code, style="LLVM") -def generate_and_write_f90_interface(build_path: str, plugin: CffiPlugin): +def generate_and_write_f90_interface(build_path: Path, plugin: CffiPlugin): generated_code = F90InterfaceGenerator.apply(plugin) write_string(generated_code, build_path, f"{plugin.name}.f90") diff --git a/tools/src/icon4pytools/py2f/typing_utils.py b/tools/src/icon4pytools/py2f/typing_utils.py index 72ba5fa286..20960e9940 100644 --- a/tools/src/icon4pytools/py2f/typing_utils.py +++ b/tools/src/icon4pytools/py2f/typing_utils.py @@ -12,16 +12,16 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.common import Dimension -from gt4py.next.type_system.type_specifications import ScalarType +from gt4py.next.type_system.type_specifications import FieldType, ScalarKind, ScalarType from gt4py.next.type_system.type_translation import from_type_hint -def parse_annotation(annotation) -> tuple[list[Dimension], ScalarType]: +def parse_annotation(annotation) -> tuple[list[Dimension], ScalarKind]: type_spec = from_type_hint(annotation) + if isinstance(type_spec, ScalarType): - dtype = type_spec.kind - dims = [] + return [], type_spec.kind + elif isinstance(type_spec, FieldType): + return type_spec.dims, type_spec.dtype.kind else: - dtype = type_spec.dtype.kind - dims = type_spec.dims - return dims, dtype + raise ValueError(f"Unsupported type specification: {type_spec}") diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index 89c2dbb59e..d1b4c24243 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -11,6 +11,10 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # flake8: noqa + +# We use gt4py type annotations and thus need to ignore this in MyPy +# mypy: disable-error-code="valid-type" + """ Wrapper module for diffusion granule. @@ -18,9 +22,8 @@ Fortran granule interfaces: - all arguments needed from external sources are passed. - passing of scalar types or fields of simple types - - """ + import numpy as np from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 diff --git a/tools/tests/f2ser/fortran_samples/no_deps_subroutine_example.f90 b/tools/tests/f2ser/fortran_samples/no_deps_subroutine_example.f90 index 021cb4aa76..cdaf8a904f 100644 --- a/tools/tests/f2ser/fortran_samples/no_deps_subroutine_example.f90 +++ b/tools/tests/f2ser/fortran_samples/no_deps_subroutine_example.f90 @@ -4,6 +4,8 @@ MODULE no_deps_example_subroutines PUBLIC :: no_deps_init, no_deps_run PRIVATE + USE foo + CONTAINS SUBROUTINE no_deps_init(a, b, c) diff --git a/tools/tests/f2ser/test_f2ser_codegen.py b/tools/tests/f2ser/test_f2ser_codegen.py index 5ca607eff9..b3a41c2972 100644 --- a/tools/tests/f2ser/test_f2ser_codegen.py +++ b/tools/tests/f2ser/test_f2ser_codegen.py @@ -30,7 +30,7 @@ def test_deserialiser_diffusion_codegen(diffusion_granule, diffusion_granule_dep def expected_no_deps_serialization_directives(): serialization_directives = [ GeneratedCode( - startln=12, + startln=14, source="\n" ' !$ser init directory="." prefix="f2ser"\n' "\n" @@ -45,7 +45,7 @@ def expected_no_deps_serialization_directives(): " !$ser data b=b", ), GeneratedCode( - startln=14, + startln=16, source="\n" " !$ser savepoint no_deps_init_out\n" "\n" @@ -58,7 +58,7 @@ def expected_no_deps_serialization_directives(): " !$ser data b=b", ), GeneratedCode( - startln=20, + startln=22, source="\n" " !$ser savepoint no_deps_run_in\n" "\n" @@ -71,7 +71,7 @@ def expected_no_deps_serialization_directives(): " !$ser data b=b", ), GeneratedCode( - startln=22, + startln=24, source="\n" " !$ser savepoint no_deps_run_out\n" "\n" From 9e2fbb410090cd84aa965b420aeb3d59525a679c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 26 Sep 2023 16:20:33 +0200 Subject: [PATCH 100/105] Ignore delete statements in fused mode (#266) This PR allows to use the DELETE statements only for the fused stencil mode. --------- Co-authored-by: Samuel --- tools/README.md | 70 ++++++ .../icon4pytools/liskov/parsing/transform.py | 5 + tools/tests/liskov/test_transform.py | 201 ++++++++++++++++++ 3 files changed, 276 insertions(+) create mode 100644 tools/tests/liskov/test_transform.py diff --git a/tools/README.md b/tools/README.md index 41f040e0dc..2a2d5c29c6 100644 --- a/tools/README.md +++ b/tools/README.md @@ -192,6 +192,66 @@ Additionally, there are the following keyword arguments: - `noaccenddata`: Takes a boolean string input and controls whether a `!$ACC END DATA` directive is generated or not. Defaults to false.

+#### `!$DSL FUSED START STENCIL()` + +This directive denotes the start of a fused stencil. Required arguments are `name`, `vertical_lower`, `vertical_upper`, `horizontal_lower`, `horizontal_upper`. The value for `name` must correspond to a stencil found in one of the stencil modules inside `icon4py`, and all fields defined in the directive must correspond to the fields defined in the respective icon4py stencil. Optionally, absolute and relative tolerances for the output fields can also be set using the `_tol` or `_abs` suffixes respectively. For each stencil, an ACC ENTER/EXIT DATA statements will be created. This ACC ENTER/EXIT DATA region contains the before fileds of the according stencil. An example call looks like this: + +```fortran + !$DSL START FUSED STENCIL(name=calculate_diagnostic_quantities_for_turbulence; & + !$DSL kh_smag_ec=kh_smag_ec(:,:,1); vn=p_nh_prog%vn(:,:,1); e_bln_c_s=p_int%e_bln_c_s(:,:,1); & + !$DSL geofac_div=p_int%geofac_div(:,:,1); diff_multfac_smag=diff_multfac_smag(:); & + !$DSL wgtfac_c=p_nh_metrics%wgtfac_c(:,:,1); div_ic=p_nh_diag%div_ic(:,:,1); & + !$DSL hdef_ic=p_nh_diag%hdef_ic(:,:,1); & + !$DSL div_ic_abs_tol=1e-18_wp; vertical_lower=2; & + !$DSL vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) +``` + +#### `!$DSL END FUSED STENCIL()` + +This directive denotes the end of a fused stencil. The required argument is `name`, which must match the name of the preceding `START STENCIL` directive. + +Note that each `START STENCIL` and `END STENCIL` will be transformed into a `DELETE` section, when using the `--fused` mode. +Together, the `START FUSED STENCIL` and `END FUSED STENCIL` directives result in the following generated code at the start and end of a stencil respectively. + +```fortran + !$ACC DATA CREATE( & + !$ACC kh_smag_e_before, & + !$ACC kh_smag_ec_before, & + !$ACC z_nabla2_e_before ) & + !$ACC IF ( i_am_accel_node ) + +#ifdef __DSL_VERIFY + !$ACC KERNELS IF( i_am_accel_node ) DEFAULT(PRESENT) ASYNC(1) + kh_smag_e_before(:, :, :) = kh_smag_e(:, :, :) + kh_smag_ec_before(:, :, :) = kh_smag_ec(:, :, :) + z_nabla2_e_before(:, :, :) = z_nabla2_e(:, :, :) + !$ACC END KERNELS +``` + +```fortran +call wrap_run_calculate_diagnostic_quantities_for_turbulence( & + kh_smag_ec=kh_smag_ec(:, :, 1), & + vn=p_nh_prog%vn(:, :, 1), & + e_bln_c_s=p_int%e_bln_c_s(:, :, 1), & + geofac_div=p_int%geofac_div(:, :, 1), & + diff_multfac_smag=diff_multfac_smag(:), & + wgtfac_c=p_nh_metrics%wgtfac_c(:, :, 1), & + div_ic=p_nh_diag%div_ic(:, :, 1), & + div_ic_before=div_ic_before(:, :, 1), & + hdef_ic=p_nh_diag%hdef_ic(:, :, 1), & + hdef_ic_before=hdef_ic_before(:, :, 1), & + div_ic_abs_tol=1e-18_wp, & + vertical_lower=2, & + vertical_upper=nlev, & + horizontal_lower=i_startidx, & + horizontal_upper=i_endidx) + +!$ACC EXIT DATA DELETE( & +!$ACC div_ic_before, & +!$ACC hdef_ic_before ) & +!$ACC IF ( i_am_accel_node ) +``` + #### `!$DSL INSERT()` This directive allows the user to generate any text that is placed between the parentheses. This is useful for situations where custom code generation is necessary. @@ -204,6 +264,16 @@ This directive allows generating an nvtx start profile data statement, and takes This directive allows generating an nvtx end profile statement. +#### `!$DSL START DELETE + +This directive allows to disable code. The code is only disabled if both the fused mode and the substition mode are enabled. +The `START DELETE` indicates the starting line from which on code is deleted. + +#### `!$DSL END DELETE` + +This directive allows to disable code. The code is only disabled if both the fused mode and the substition mode are enabled. +The `END DELETE` indicates the ending line from which on code is deleted. + #### `!$DSL ENDIF()` This directive generates an `#endif` statement. diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 3537a2dc63..f5ace9add4 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -54,6 +54,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: else: logger.info("Removing fused stencils.") self._remove_fused_stencils() + self._remove_delete() return self.parsed @@ -113,3 +114,7 @@ def _remove_stencils(self, stencils_to_remove: list[CodeGenInput]) -> None: def _remove_fused_stencils(self) -> None: self.parsed.StartFusedStencil = [] self.parsed.EndFusedStencil = [] + + def _remove_delete(self) -> None: + self.parsed.StartDelete = [] + self.parsed.EndDelete = [] diff --git a/tools/tests/liskov/test_transform.py b/tools/tests/liskov/test_transform.py new file mode 100644 index 0000000000..4c2a60454c --- /dev/null +++ b/tools/tests/liskov/test_transform.py @@ -0,0 +1,201 @@ +# 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 + + +import pytest + +from icon4pytools.liskov.codegen.integration.interface import ( + BoundsData, + DeclareData, + EndCreateData, + EndDeleteData, + EndFusedStencilData, + EndIfData, + EndProfileData, + EndStencilData, + FieldAssociationData, + ImportsData, + InsertData, + IntegrationCodeInterface, + StartCreateData, + StartDeleteData, + StartFusedStencilData, + StartProfileData, + StartStencilData, +) +from icon4pytools.liskov.parsing.transform import StencilTransformer + + +@pytest.fixture +def integration_code_interface(): + start_fused_stencil_data = StartFusedStencilData( + name="fused_stencil1", + fields=[ + FieldAssociationData("scalar1", "scalar1", inp=True, out=False, dims=None), + FieldAssociationData("inp1", "inp1(:,:,1)", inp=True, out=False, dims=2), + FieldAssociationData("out1", "out1(:,:,1)", inp=False, out=True, dims=2, abs_tol="0.5"), + FieldAssociationData( + "out2", + "p_nh%prog(nnew)%out2(:,:,1)", + inp=False, + out=True, + dims=3, + abs_tol="0.2", + ), + FieldAssociationData("out3", "p_nh%prog(nnew)%w(:,:,jb)", inp=False, out=True, dims=2), + FieldAssociationData("out4", "p_nh%prog(nnew)%w(:,:,1,2)", inp=False, out=True, dims=3), + FieldAssociationData( + "out5", "p_nh%prog(nnew)%w(:,:,:,ntnd)", inp=False, out=True, dims=3 + ), + FieldAssociationData( + "out6", "p_nh%prog(nnew)%w(:,:,1,ntnd)", inp=False, out=True, dims=3 + ), + ], + bounds=BoundsData("1", "10", "-1", "-10"), + startln=1, + acc_present=False, + ) + end_fused_stencil_data = EndFusedStencilData(name="stencil1", startln=4) + start_stencil_data1 = StartStencilData( + name="stencil1", + fields=[ + FieldAssociationData("scalar1", "scalar1", inp=True, out=False, dims=None), + FieldAssociationData("inp1", "inp1(:,:,1)", inp=True, out=False, dims=2), + FieldAssociationData("out1", "out1(:,:,1)", inp=False, out=True, dims=2, abs_tol="0.5"), + FieldAssociationData( + "out2", + "p_nh%prog(nnew)%out2(:,:,1)", + inp=False, + out=True, + dims=3, + abs_tol="0.2", + ), + FieldAssociationData("out3", "p_nh%prog(nnew)%w(:,:,jb)", inp=False, out=True, dims=2), + FieldAssociationData("out4", "p_nh%prog(nnew)%w(:,:,1,2)", inp=False, out=True, dims=3), + FieldAssociationData( + "out5", "p_nh%prog(nnew)%w(:,:,:,ntnd)", inp=False, out=True, dims=3 + ), + FieldAssociationData( + "out6", "p_nh%prog(nnew)%w(:,:,1,ntnd)", inp=False, out=True, dims=3 + ), + ], + bounds=BoundsData("1", "10", "-1", "-10"), + startln=2, + acc_present=False, + mergecopy=False, + copies=True, + ) + end_stencil_data1 = EndStencilData( + name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False + ) + start_stencil_data2 = StartStencilData( + name="stencil2", + fields=[ + FieldAssociationData("scalar1", "scalar1", inp=True, out=False, dims=None), + FieldAssociationData("inp1", "inp1(:,:,1)", inp=True, out=False, dims=2), + FieldAssociationData("out1", "out1(:,:,1)", inp=False, out=True, dims=2, abs_tol="0.5"), + FieldAssociationData( + "out2", + "p_nh%prog(nnew)%out2(:,:,1)", + inp=False, + out=True, + dims=3, + abs_tol="0.2", + ), + FieldAssociationData("out3", "p_nh%prog(nnew)%w(:,:,jb)", inp=False, out=True, dims=2), + FieldAssociationData("out4", "p_nh%prog(nnew)%w(:,:,1,2)", inp=False, out=True, dims=3), + FieldAssociationData( + "out5", "p_nh%prog(nnew)%w(:,:,:,ntnd)", inp=False, out=True, dims=3 + ), + FieldAssociationData( + "out6", "p_nh%prog(nnew)%w(:,:,1,ntnd)", inp=False, out=True, dims=3 + ), + ], + bounds=BoundsData("1", "10", "-1", "-10"), + startln=5, + acc_present=False, + mergecopy=False, + copies=True, + ) + end_stencil_data2 = EndStencilData( + name="stencil2", startln=6, noendif=False, noprofile=False, noaccenddata=False + ) + declare_data = DeclareData( + startln=7, + declarations={"field2": "(nproma, p_patch%nlev, p_patch%nblks_e)"}, + ident_type="REAL(wp)", + suffix="before", + ) + imports_data = ImportsData(startln=8) + start_create_data = StartCreateData(extra_fields=["foo", "bar"], startln=9) + end_create_data = EndCreateData(startln=11) + endif_data = EndIfData(startln=12) + start_profile_data = StartProfileData(startln=13, name="test_stencil") + end_profile_data = EndProfileData(startln=14) + insert_data = InsertData(startln=15, content="print *, 'Hello, World!'") + start_delete_data = StartDeleteData(startln=16) + end_delete_data = EndDeleteData(startln=17) + + return IntegrationCodeInterface( + StartStencil=[start_stencil_data1, start_stencil_data2], + EndStencil=[end_stencil_data1, end_stencil_data2], + StartFusedStencil=[start_fused_stencil_data], + EndFusedStencil=[end_fused_stencil_data], + StartDelete=[start_delete_data], + EndDelete=[end_delete_data], + Declare=[declare_data], + Imports=imports_data, + StartCreate=[start_create_data], + EndCreate=[end_create_data], + EndIf=[endif_data], + StartProfile=[start_profile_data], + EndProfile=[end_profile_data], + Insert=[insert_data], + ) + + +@pytest.fixture +def stencil_transform_fused(integration_code_interface): + return StencilTransformer(integration_code_interface, fused=True) + + +@pytest.fixture +def stencil_transform_unfused(integration_code_interface): + return StencilTransformer(integration_code_interface, fused=False) + + +def test_transform_fused( + stencil_transform_fused, +): + # Check that the transformed interface is as expected + transformed = stencil_transform_fused() + assert len(transformed.StartFusedStencil) == 1 + assert len(transformed.EndFusedStencil) == 1 + assert len(transformed.StartStencil) == 1 + assert len(transformed.EndStencil) == 1 + assert len(transformed.StartDelete) == 2 + assert len(transformed.EndDelete) == 2 + + +def test_transform_unfused( + stencil_transform_unfused, +): + # Check that the transformed interface is as expected + transformed = stencil_transform_unfused() + + assert not transformed.StartFusedStencil + assert not transformed.EndFusedStencil + assert len(transformed.StartStencil) == 2 + assert len(transformed.EndStencil) == 2 + assert not transformed.StartDelete + assert not transformed.EndDelete From ad6b8b7d3919e8fedaf35a21b68a88e705cd319e Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 26 Sep 2023 18:51:37 +0200 Subject: [PATCH 101/105] Add custom marker slow_tests to skip 4 slow stencils for testing --- .../advection_tests/test_divide_flux_area_list_stencil_01.py | 4 +++- .../test_prep_gauss_quadrature_c_list_stencil.py | 2 ++ .../advection_tests/test_prep_gauss_quadrature_c_stencil.py | 2 ++ .../advection_tests/test_recon_lsq_cell_c_svd_stencil.py | 2 ++ model/atmosphere/advection/pyproject.toml | 3 ++- model/tox.ini | 4 ++-- pytest.ini | 2 ++ tox.ini | 4 ++-- 8 files changed, 17 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py index 82cba3671a..a2a2b0f3fc 100644 --- a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py @@ -12,6 +12,8 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest + from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider @@ -658,7 +660,7 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask, ) - +@pytest.mark.slow_tests def test_divide_flux_area_list_stencil_01(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py index 1bdc3b7144..202669e377 100644 --- a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( @@ -315,6 +316,7 @@ def prep_gauss_quadrature_c_list_stencil_numpy( ) +@pytest.mark.slow_tests def test_prep_gauss_quadrature_c_list_stencil(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py index a58cda650e..876869dfbb 100644 --- a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( prep_gauss_quadrature_c_stencil, @@ -283,6 +284,7 @@ def prep_gauss_quadrature_c_stencil_numpy( ) +@pytest.mark.slow_tests def test_prep_gauss_quadrature_c_stencil(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py index 74163d0ff2..7dda795384 100644 --- a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.model.atmosphere.advection.recon_lsq_cell_c_svd_stencil import ( @@ -206,6 +207,7 @@ def recon_lsq_cell_c_svd_stencil_numpy( ) +@pytest.mark.slow_tests def test_recon_lsq_cell_c_svd_stencil(): mesh = SimpleMesh() p_cc = random_field(mesh, CellDim, KDim) diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index 65d527bff9..ffe43aff5f 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -110,7 +110,8 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'tests' +testpaths = ['tests', 'advection_tests'] +markers = 'slow_tests: marks tests as slow' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} diff --git a/model/tox.ini b/model/tox.ini index 09c7bbca4a..d4a9156378 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src atmosphere/advection/src common/src driver/src - pytest -v -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -m "not slow_tests" -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src atmosphere/advection/src common/src driver/src + pytest -v -s -n auto -m "not slow_tests" --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html diff --git a/pytest.ini b/pytest.ini index 84a9d34b4e..4c4723375b 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,5 @@ [pytest] norecursedirs = _external_src python_functions = test_* bench_* +markers = + slow_tests: mark test as slow. diff --git a/tox.ini b/tox.ini index d32caf0541..622d25dea1 100644 --- a/tox.ini +++ b/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/atmosphere/advection/src model/driver/src model/common/src tools/src - pytest -v -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -m "not slow_tests" -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/atmosphere/advection/src model/driver/src model/common/src tools/src + pytest -v -s -n auto -m "not slow_tests" --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html From a50b1af9338ee085083ee327e0d07a45b18acd67 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 26 Sep 2023 18:54:31 +0200 Subject: [PATCH 102/105] Run pre-commit --- .../advection_tests/test_divide_flux_area_list_stencil_01.py | 2 +- model/atmosphere/advection/pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py index a2a2b0f3fc..99e25a98cb 100644 --- a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py @@ -13,7 +13,6 @@ import numpy as np import pytest - from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider @@ -660,6 +659,7 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask, ) + @pytest.mark.slow_tests def test_divide_flux_area_list_stencil_01(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index ffe43aff5f..f26f25f3b4 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -110,8 +110,8 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = ['tests', 'advection_tests'] markers = 'slow_tests: marks tests as slow' +testpaths = ['tests', 'advection_tests'] [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} From 24a6788081367a4f290924fb8ffe3ee1a5461475 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 26 Sep 2023 18:51:37 +0200 Subject: [PATCH 103/105] Add custom marker slow_tests to skip 4 slow stencils for testing --- .../advection_tests/test_divide_flux_area_list_stencil_01.py | 3 +++ model/atmosphere/advection/pyproject.toml | 1 + 2 files changed, 4 insertions(+) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py index 99e25a98cb..ef2fd50497 100644 --- a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py @@ -659,7 +659,10 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask, ) +<<<<<<< HEAD +======= +>>>>>>> 168f9fa... Add custom marker slow_tests to skip 4 slow stencils for testing @pytest.mark.slow_tests def test_divide_flux_area_list_stencil_01(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index f26f25f3b4..27e7f6db87 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -112,6 +112,7 @@ warn_unused_ignores = true [tool.pytest.ini_options] markers = 'slow_tests: marks tests as slow' testpaths = ['tests', 'advection_tests'] +markers = 'slow_tests: marks tests as slow' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} From d4436c0b864f00fc74fbdb02de43350ca4cd6e0f Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 26 Sep 2023 18:54:31 +0200 Subject: [PATCH 104/105] Run pre-commit --- .../advection_tests/test_divide_flux_area_list_stencil_01.py | 3 --- model/atmosphere/advection/pyproject.toml | 1 - 2 files changed, 4 deletions(-) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py index ef2fd50497..99e25a98cb 100644 --- a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py @@ -659,10 +659,7 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch2_4_lat_vmask, ) -<<<<<<< HEAD -======= ->>>>>>> 168f9fa... Add custom marker slow_tests to skip 4 slow stencils for testing @pytest.mark.slow_tests def test_divide_flux_area_list_stencil_01(): mesh = SimpleMesh() diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index 27e7f6db87..f26f25f3b4 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -112,7 +112,6 @@ warn_unused_ignores = true [tool.pytest.ini_options] markers = 'slow_tests: marks tests as slow' testpaths = ['tests', 'advection_tests'] -markers = 'slow_tests: marks tests as slow' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} From aad80b885746b3245cbe80367998c1b6664b43f3 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 28 Sep 2023 13:23:49 +0200 Subject: [PATCH 105/105] pre-commit fix --- .../model/common/test_utils/simple_mesh.py | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index ee7a699cb1..763616cc7f 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -17,6 +17,7 @@ from gt4py.next.iterator.embedded import NeighborTableOffsetProvider, StridedNeighborOffsetProvider from icon4py.model.common.dimension import ( + C2E2C2E2CDim, C2E2CDim, C2E2CODim, C2EDim, @@ -34,8 +35,6 @@ V2CDim, V2EDim, VertexDim, - C2E2C2E2CDim, - ECDim ) @@ -375,24 +374,24 @@ class SimpleMeshData: c2e2c2e2c_table = np.asarray( [ - [15, 4, 3, 12, 14, 1, 7, 6, 2], # 1c - [16, 5, 4, 12, 13, 2, 8, 7, 0], - [17, 3, 5, 13, 14, 0, 6, 8, 1], - [0, 6, 2, 17, 5, 9, 10, 15, 4], - [1, 7, 0, 15, 3, 16, 5, 10, 11], # 5c - [2, 8, 1, 4, 16, 17, 3, 9, 11], - [3, 10, 9, 2, 0, 7, 13, 8, 12], - [4, 11, 10, 0, 1, 8, 14, 6, 13], - [5, 9, 11, 1, 2, 3, 12, 7, 14], - [6, 12, 8, 5, 11, 3, 10, 16, 15], # 10c - [7, 13, 6, 3, 9, 4, 11, 16, 17], - [8, 14, 7, 4, 10, 5, 9, 15, 17], - [9, 16, 15, 8, 6, 1, 13, 0, 14], - [10, 17, 16, 6, 7, 2, 14, 1, 12], - [11, 15, 17, 7, 8, 2, 13, 0, 12], # 15c - [12, 0, 14, 11, 17, 9, 16, 3, 4], - [13, 1, 12, 9, 15, 10, 17, 4, 5], - [14, 2, 13, 10, 16, 5, 3, 11, 15], + [15, 4, 3, 12, 14, 1, 7, 6, 2], # 1c + [16, 5, 4, 12, 13, 2, 8, 7, 0], + [17, 3, 5, 13, 14, 0, 6, 8, 1], + [0, 6, 2, 17, 5, 9, 10, 15, 4], + [1, 7, 0, 15, 3, 16, 5, 10, 11], # 5c + [2, 8, 1, 4, 16, 17, 3, 9, 11], + [3, 10, 9, 2, 0, 7, 13, 8, 12], + [4, 11, 10, 0, 1, 8, 14, 6, 13], + [5, 9, 11, 1, 2, 3, 12, 7, 14], + [6, 12, 8, 5, 11, 3, 10, 16, 15], # 10c + [7, 13, 6, 3, 9, 4, 11, 16, 17], + [8, 14, 7, 4, 10, 5, 9, 15, 17], + [9, 16, 15, 8, 6, 1, 13, 0, 14], + [10, 17, 16, 6, 7, 2, 14, 1, 12], + [11, 15, 17, 7, 8, 2, 13, 0, 12], # 15c + [12, 0, 14, 11, 17, 9, 16, 3, 4], + [13, 1, 12, 9, 15, 10, 17, 4, 5], + [14, 2, 13, 10, 16, 5, 3, 11, 15], ] )