diff --git a/CODING_GUIDELINES.md b/CODING_GUIDELINES.md new file mode 100644 index 0000000000..8631c576b6 --- /dev/null +++ b/CODING_GUIDELINES.md @@ -0,0 +1,168 @@ +# Coding Guidelines + +## Software design + +Designing a clean architecture for a software product is a difficult task and developers can easily overlook it under pressure or when they are eager to deliver exciting new features. However, choosing the quick solution instead of a clean one will likely lead to a large amount of additional work in the future (_technical debt_). + +To keep technical debt at acceptable levels, design and implementation of new features should follow best practices: + +1. Make sure your code is always covered by automatic testing to ensure its correctness. Where unit tests are impractical, use integration tests. +2. Adhere to the [SOLID](https://en.wikipedia.org/wiki/SOLID) principles of software design. +3. Do not reinvent the wheel: if someone solved your problem within the project or in a third party library, consider using their solution or extending it before writing a new component for the same purpose. +4. _You ain't gonna need it_ ([YAGNI](https://en.wikipedia.org/wiki/You_aren%27t_gonna_need_it)): do not design solutions for problems that might come up in the future, as chances are that you will never need that code. Focus on current problems and prepare for future requirements by writing clean code. +5. _Do not repeat yourself_ ([DRY](https://en.wikipedia.org/wiki/Don%27t_repeat_yourself)): if you are writing the same code snippet in several places, extract it into a function. +6. Use meaningful names: the purpose of an object should be clear from its name. Usually, class names are nouns and function names are verbs. + + +## Code Style + +We follow the [Google Python Style Guide][google-style-guide] with a few minor changes (mentioned below). Since the best way to remember something is to understand the reasons behind it, make sure you go through the style guide at least once, paying special attention to the discussions in the _Pros_, _Cons_, and _Decision_ subsections. + +We deviate from the [Google Python Style Guide][google-style-guide] only in the following points: + +- We use [`ruff-linter`][ruff-linter] instead of [`pylint`][pylint]. +- We use [`ruff-formatter`][ruff-formatter] for source code and imports formatting, which may work differently than indicated by the guidelines in section [_3. Python Style Rules_](https://google.github.io/styleguide/pyguide.html#3-python-style-rules). For example, maximum line length is set to 100 instead of 79 (although docstring lines should still be limited to 79). +- According to subsection [_2.19 Power Features_](https://google.github.io/styleguide/pyguide.html#219-power-features), direct use of _power features_ (e.g. custom metaclasses, import hacks, reflection) should be avoided, but standard library classes that internally use these power features are accepted. Following the same spirit, we allow the use of power features in infrastructure code with similar functionality and scope as the Python standard library. +- According to subsection [_3.19.12 Imports For Typing_](https://google.github.io/styleguide/pyguide.html#31912-imports-for-typing), symbols from `typing` and `collections.abc` modules used in type annotations _"can be imported directly to keep common annotations concise and match standard typing practices"_. Following the same spirit, we allow symbols to be imported directly from third-party or internal modules when they only contain a collection of frequently used typying definitions. + +Further guidance and repository- or package specific guidelines can be found in the respective `docs` folders and `README.md` files. +In particular +- [naming conventions for stencils](model/docs/stencil_naming_convention.md) +- [liskov integration style guide ](tools/docs/ICON_Liskov_integration_style_guide.md) + +### Common questions + +- `pass` vs `...` (`Ellipsis`) + + `pass` is the _no-op_ statement in Python and `...` is a literal value (called _Ellipsis_) introduced for slicing collections of unknown number of dimensions. Although they are very different in nature, both of them are used in places where a statement is required purely for syntactic reasons, and there is not yet a clear standard practice in the community about when to use one or the other. We decided to align with the common pattern of using `...` in the body of empty function definitions working as placeholders for actual implementations defined somewhere else (e.g. type stubs, abstract methods and methods appearing in `Protocol` classes) and `pass` in any other place where its usage is mixed with actual statements. + + ```python + # Correct use of `...` as the empty body of an abstract method + class AbstractFoo: + @abstractmethod + def bar(self) -> Bar: + ... + + # Correct use of `pass` when mixed with other statements + try: + resource.load(id=42) + except ResourceException: + pass + ``` + + +### Error messages + +Error messages should be written as sentences, starting with a capital letter and ending with a period (avoid exclamation marks). Try to be informative without being verbose. +The message should be kept to one sentence if reasonably possible. +Code objects such as 'ClassNames' and 'function_names' should be enclosed in single quotes, and so should string values used for message interpolation. + +Examples: + +```python +raise ValueError(f"Invalid argument 'dimension': should be of type 'Dimension', got '{dimension.type}'.") +``` + +Interpolated integer values do not need double quotes, if they are indicating an amount. Example: + +```python +raise ValueError(f"Invalid number of arguments: expected 3 arguments, got {len(args)}.") +``` + +The double quotes can also be dropped when presenting a sequence of values. In this case the message should be rephrased so the sequence is separated from the text by a colon ':'. + +```python +raise ValueError(f"unexpected keyword arguments: {', '.join(set(kwarg_names} - set(expected_kwarg_names)))}.") +``` + +### Docstrings +We encourage to add doc strings for functions, classes and modules if they help the reader understand the code and contain information that is not obvious from the code itself. While we do not yet generate API documentation from doc strings we might do so in the future using [Sphinx][sphinx] and some extensions such as [Sphinx-autodoc][sphinx-autodoc] and [Sphinx-napoleon][sphinx-napoleon]. These follow the Google Python Style Guide docstring conventions to automatically format the generated documentation. A complete overview can be found here: [Example Google Style Python Docstrings](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html#example-google). + +Sphinx supports the [reStructuredText][sphinx-rest] (reST) markup language for defining additional formatting options in the generated documentation, however section [_3.8 Comments and Docstrings_](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) of the Google Python Style Guide does not specify how to use markups in docstrings. As a result, we decided to forbid reST markup in docstrings, except for the following cases: + +- Cross-referencing other objects using Sphinx text roles for the [Python domain](https://www.sphinx-doc.org/en/master/usage/restructuredtext/domains.html#the-python-domain) (as explained [here](https://www.sphinx-doc.org/en/master/usage/restructuredtext/domains.html#python-roles)). +- Very basic formatting markup to improve _readability_ of the generated documentation without obscuring the source docstring (e.g. ` ``literal`` ` strings, bulleted lists). + +We highly encourage the [doctest][doctest] format for code examples in docstrings. In fact, doctest runs code examples and makes sure they are in sync with the codebase. + +### Module structure + +In general, you should structure new Python modules in the following way: + +1. _shebang_ line: `#! /usr/bin/env python3` (only for **executable scripts**!). +2. License header (see `.license_header.txt`, it is added automatically by `pre-commit` hook) +3. Module docstring. +4. Imports, alphabetically ordered within each block (fixed automatically by `ruff-formatter`): + 1. Block of imports from the standard library. + 2. Block of imports from general third party libraries using standard shortcuts when customary (e.g. `numpy as np`). + 3. Block of imports from specific modules of the project. +5. Definition of exported symbols (optional, mainly for re-exporting symbols from other modules): + +```python +__all__ = ["func_a", "CONST_B"] +``` + +6. Public constants and typing definitions. +7. Module contents organized in a convenient way for understanding how the pieces of code fit together, usually defining functions before classes. + +Try to keep sections and items logically ordered, add section separator comments to make section boundaries explicit when needed. If there is not a single evident logical order, pick the order you consider best or use alphabetical order. + +Consider configuration files as another type of source code and apply the same criteria, using comments when possible for better readability. + +### Ignoring QA errors + +You may occasionally need to disable checks from _quality assurance_ (QA) tools (e.g. linters, type checkers, etc.) on specific lines as some tool might not be able to fully understand why a certain piece of code is needed. This is usually done with special comments, e.g. `# noqa: F401`, `# type: ignore`. However, you should **only** ignore QA errors when you fully understand their source and rewriting your code to pass QA checks would make it less readable. Additionally, you should add a short descriptive code if possible (check [ruff rules][ruff-rules] and [mypy error codes][mypy-error-codes] for reference): + +```python +f = lambda: 'empty' # noqa: E731 [lambda-assignment] +``` + +and, if needed, a brief comment for future reference: + +```python +... +return undeclared_symbol # noqa: F821 [undefined-name] on purpose to trigger black-magic +``` + +## Testing + +Testing components is a critical part of a software development project. We follow standard practices in software development and write unit, integration, and regression tests. Note that even though [doctests][doctest] are great for documentation purposes, they lack many features and are difficult to debug. Hence, they should not be used as replacement for proper unit tests except in trivial cases. + +### Test suite folder structure +In each package tests are organized under the `tests` folder. The `tests` folder should not be a package but the contained test suites may be python packages. + +Test suites in folders `stencil_tests` are generally run in integration mode with [icon-exclaim](https://github.com/C2SM/icon-exclaim) and should only contain tests for the GT4Py stencils that might be integrated into ICON. + + +## Further + + + + + + +[doctest]: https://docs.python.org/3/library/doctest.html +[google-style-guide]: https://google.github.io/styleguide/pyguide.html +[mypy]: https://mypy.readthedocs.io/ +[mypy-error-codes]: https://mypy.readthedocs.io/en/stable/error_code_list.html +[pre-commit]: https://pre-commit.com/ +[pylint]: https://pylint.pycqa.org/ +[ruff-formatter]: https://docs.astral.sh/ruff/formatter/ +[ruff-linter]: https://docs.astral.sh/ruff/linter/ +[ruff-rules]: https://docs.astral.sh/ruff/rules/ +[sphinx]: https://www.sphinx-doc.org +[sphinx-autodoc]: https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html +[sphinx-napoleon]: https://sphinxcontrib-napoleon.readthedocs.io/en/latest/index.html# +[sphinx-rest]: https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html diff --git a/README.md b/README.md index 2b4738ba6d..0535887e69 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ python3.10 -m venv .venv # Activate the virtual environment and make sure that 'wheel' is installed source .venv/bin/activate -pip install --upgrade wheel +pip install --upgrade wheel pip setuptools # Install all the ICON4Py packages and its dependencies # External dependencies would be checked out at './_external_src' diff --git a/ci/default.yml b/ci/default.yml index 4a85316030..7c17844fac 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -17,7 +17,7 @@ test_model_stencils: # exclude slow test configurations - if: $BACKEND == "roundtrip" && $GRID == "icon_grid" when: never - - when: always + - when: on_success test_tools: extends: .test_template diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index 532ed94966..7e0312f5a2 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -23,7 +23,8 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py-common>=0.0.5" + "icon4py-common>=0.0.5", + "nanobind<2.0.0" ] description = "ICON advection." dynamic = ['version'] 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 3537d58045..0f0d962e1a 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 @@ -27,9 +27,7 @@ def _btraj_dreg_stencil_02( lvn_pos = where(p_vn >= 0.0, True, False) 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) - ) + opt_famask_dsl = where(traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), 1, 0) return opt_famask_dsl 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 d2c2fd2c72..381c1e31a9 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 @@ -50,20 +50,37 @@ def _btraj_dreg_stencil_03( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: + # logical switch for MERGE operations: True for p_vn >= 0 lvn_pos = where(p_vn >= 0.0, True, False) + # get line and block indices of upwind cell 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_rel_idx_dsl = where(lvn_pos, 0, 1) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) + # departure points of the departure cell. Point 1 belongs to edge-vertex 1, + # point 2 belongs to edge_vertex 2. + # + # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in normal direction depart_pts_1_x = edge_verts_1_x - p_vn * p_dt + # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in tangential direction depart_pts_1_y = edge_verts_1_y - p_vt * p_dt + # position of vertex 3 in normal direction depart_pts_2_x = edge_verts_2_x - p_vn * p_dt + # position of vertex 3 in tangential direction depart_pts_2_y = edge_verts_2_y - p_vt * p_dt + # determine correct position on tangential plane 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) + # Calculate position of departure region vertices in a translated + # coordinate system. The origin is located at the circumcenter + # of the upwind cell. The distance vectors point from the cell center + # to the vertices. + + # Take care of correct counterclockwise numbering below + 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 @@ -73,15 +90,24 @@ def _btraj_dreg_stencil_03( 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 + # In a last step, these distance vectors are transformed into a rotated + # geographical coordinate system, which still has its origin at the circumcenter + # of the upwind cell. Now the coordinate axes point to local East and local + # North. + + # Determine primal and dual normals of the cell lying in the direction of vn 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])) + # components in longitudinal direction 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 + + # components in latitudinal direction 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 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 2fa1051c61..242e43d6ef 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 @@ -107,8 +107,15 @@ 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) + # avoid division with zero + d1 = line1_p2_lon - line1_p1_lon + d1 = where(d1 != 0.0, d1, line1_p2_lon) + + d2 = line2_p2_lon - line2_p1_lon + d2 = where(d2 != 0.0, d2, line2_p2_lon) + + m1 = (line1_p2_lat - line1_p1_lat) / d1 + m2 = (line2_p2_lat - line2_p1_lat) / d2 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) @@ -228,7 +235,7 @@ def _divide_flux_area_list_stencil_01( 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) + famask_bool = where(famask_int == 1, True, False) # ------------------------------------------------- Case 1 mask_case1 = lintersect_line1 & lintersect_line2 & famask_bool ps1_x, ps1_y = line_intersect( 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 ec0d9e0ad4..938f605989 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 @@ -77,7 +77,7 @@ def _divide_flux_area_list_stencil_02( Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], ]: - famask_bool = where(famask_int == int32(1), True, False) + famask_bool = where(famask_int == 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( @@ -133,22 +133,22 @@ def _divide_flux_area_list_stencil_02( patch1_cell_idx_vmask = where( famask_bool, where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), - int32(0), + 0, ) patch2_cell_idx_vmask = where( famask_bool, where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), - int32(0), + 0, ) patch1_cell_blk_vmask = where( famask_bool, where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), - int32(0), + 0, ) patch2_cell_blk_vmask = where( famask_bool, where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), - int32(0), + 0, ) return ( 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 154bb669e5..0da8565e90 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 @@ -44,61 +44,61 @@ def _hflux_ffsl_hybrid_stencil_01a( ) -> Field[[EdgeDim, KDim], float]: p_out_e_hybrid_1a = ( where( - patch0_cell_rel_idx_dsl == int32(1), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 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), + patch0_cell_rel_idx_dsl == 1, z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0]), ) 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 2380296e9e..aeb1ff5db1 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 @@ -12,41 +12,31 @@ # 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 +from gt4py.next.ffront.fbuiltins import Field, where from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator def _hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: int32, ) -> Field[[EdgeDim, KDim], float]: p_mflx_tracer_h_out = where( - refin_ctrl == bound, - p_mflx_tracer_h, - where( - p_mflx_tracer_h >= 0.0, - p_mflx_tracer_h * r_m(E2C[0]), - p_mflx_tracer_h * r_m(E2C[1]), - ), + p_mflx_tracer_h >= 0.0, + p_mflx_tracer_h * r_m(E2C[0]), + p_mflx_tracer_h * r_m(E2C[1]), ) return p_mflx_tracer_h_out @program def hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: int32, ): _hflx_limiter_pd_stencil_02( - refin_ctrl, r_m, p_mflx_tracer_h, - bound, out=p_mflx_tracer_h, ) 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 59ad061a9d..cd8069e55e 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 @@ -27,14 +27,14 @@ def _hor_adv_stencil_01( geofac_div: Field[[CEDim], float], p_dtime: float, ) -> Field[[CellDim, KDim], float]: - tracer_new = ( + tracer_new_hor = ( 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 + return tracer_new_hor @program @@ -45,7 +45,7 @@ def hor_adv_stencil_01( rhodz_now: Field[[CellDim, KDim], float], rhodz_new: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], - tracer_new: Field[[CellDim, KDim], float], + tracer_new_hor: Field[[CellDim, KDim], float], p_dtime: float, ): _hor_adv_stencil_01( @@ -56,5 +56,5 @@ def hor_adv_stencil_01( rhodz_new, geofac_div, p_dtime, - out=tracer_new, + out=tracer_new_hor, ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py deleted file mode 100644 index 10fa5b4998..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py +++ /dev/null @@ -1,28 +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, GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import broadcast - -from icon4py.model.common.dimension import CellDim - - -@field_operator -def _init_zero_c() -> Field[[CellDim], float]: - return broadcast(0.0, (CellDim,)) - - -@program(grid_type=GridType.UNSTRUCTURED) -def init_zero_c(field: Field[[CellDim], float]): - _init_zero_c(out=field) 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 941cb99661..1e01270bb6 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 @@ -33,6 +33,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_2: Field[[ECDim], wpfloat], p_dthalf: wpfloat, ) -> tuple[ + Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], vpfloat], @@ -41,6 +42,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])) + p_cell_rel_idx_dsl = where(lvn_pos, 0, 1) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( @@ -67,7 +69,13 @@ 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, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat) + return ( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + astype(p_distv_bary_1, vpfloat), + astype(p_distv_bary_2, vpfloat), + ) @program(grid_type=GridType.UNSTRUCTURED) @@ -83,6 +91,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2: Field[[ECDim], wpfloat], dual_normal_cell_2: Field[[ECDim], wpfloat], 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], vpfloat], p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat], @@ -104,7 +113,13 @@ 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_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_distv_bary_1, + p_distv_bary_2, + ), domain={ EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), 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 7bd6a261ba..6851ded2a4 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 @@ -95,7 +95,7 @@ def _prep_gauss_quadrature_c_list_stencil( 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) + famask_bool = where(famask_int == 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) 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 2df727e075..3d5e420093 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 @@ -47,61 +47,61 @@ def _upwind_hflux_miura3_stencil_01( p_out_e_miura3 = ( ( where( - cell_rel_idx_dsl == int32(1), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 1, z_lsq_coeff_10(E2C[1]), z_lsq_coeff_10(E2C[0]), ) 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 45823abb00..e8f8a15da9 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 @@ -30,19 +30,19 @@ def _upwind_hflux_miura_cycl_stencil_01( ) -> Field[[EdgeDim, KDim], float]: z_tracer_mflx_dsl = ( where( - cell_rel_idx_dsl == int32(1), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 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), + cell_rel_idx_dsl == 1, z_lsq_coeff_3_dsl(E2C[1]), z_lsq_coeff_3_dsl(E2C[0]), ) 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 1c9e1779de..6e9354ceb1 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 @@ -35,9 +35,7 @@ def _upwind_hflux_miura_cycl_stencil_02( 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 + neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) if nsub == 1 else z_rhofluxdiv_c ) z_fluxdiv_c_dsl = neighbor_sum(z_tracer_mflx(C2E) * geofac_div, axis=C2EDim) 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 3edf4fc43a..1932c189f8 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 @@ -29,11 +29,9 @@ 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])) + where(cell_rel_idx_dsl == 1, z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) + + distv_bary_1 * where(cell_rel_idx_dsl == 1, z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) + + distv_bary_2 * where(cell_rel_idx_dsl == 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/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py index e5b325d268..61c3dcfde5 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 @@ -34,7 +34,7 @@ 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, 1, 0) return l_limit 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 8717d8fd06..df91a37f19 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 @@ -27,7 +27,7 @@ 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 != 0, where( (p_cc < minimum(p_face, p_face(Koff[1]))), (p_cc, p_cc), 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 44f933b18f..70cd6b6629 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 @@ -12,13 +12,13 @@ # 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 +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator -def _vert_adv_stencil_01( +def _vert_adv_stencil_01a( tracer_now: Field[[CellDim, KDim], float], rhodz_now: Field[[CellDim, KDim], float], p_mflx_tracer_v: Field[[CellDim, KDim], float], @@ -35,6 +35,42 @@ def _vert_adv_stencil_01( return tracer_new +@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], + k: Field[[KDim], int32], + p_dtime: float, + ivadv_tracer: int32, + iadv_slev_jt: int32, +) -> Field[[CellDim, KDim], float]: + k = broadcast(k, (CellDim, KDim)) + + tracer_new = ( + where( + (iadv_slev_jt <= k), + _vert_adv_stencil_01a( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + p_dtime, + ), + tracer_now, + ) + if (ivadv_tracer != 0) + else tracer_now + ) + + return tracer_new + + @program(grid_type=GridType.UNSTRUCTURED) def vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], @@ -43,8 +79,11 @@ def vert_adv_stencil_01( deepatmo_divzl: Field[[KDim], float], deepatmo_divzu: Field[[KDim], float], rhodz_new: Field[[CellDim, KDim], float], - tracer_new: Field[[CellDim, KDim], float], + k: Field[[KDim], int32], p_dtime: float, + ivadv_tracer: int32, + iadv_slev_jt: int32, + tracer_new: Field[[CellDim, KDim], float], ): _vert_adv_stencil_01( tracer_now, @@ -53,6 +92,9 @@ def vert_adv_stencil_01( deepatmo_divzl, deepatmo_divzu, rhodz_new, + k, p_dtime, + ivadv_tracer, + iadv_slev_jt, out=tracer_new, ) diff --git a/model/atmosphere/advection/tests/stencil_tests/__init__.py b/model/atmosphere/advection/tests/advection_stencil_tests/__init__.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/__init__.py rename to model/atmosphere/advection/tests/advection_stencil_tests/__init__.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py similarity index 98% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py index d602eae212..f04e6369a8 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py @@ -42,8 +42,8 @@ def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs): opt_famask_dsl = np.where( traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), - int32(1), - int32(0), + 1, + 0, ) return dict(opt_famask_dsl=opt_famask_dsl) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py index 927bccd654..b24f870440 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py @@ -91,7 +91,7 @@ def reference( 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)) + p_cell_rel_idx_dsl = np.where(lvn_pos, 0, 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 diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py similarity index 97% rename from model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py index 392e01e4a1..76c8f5c445 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -22,6 +22,7 @@ from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, + constant_field, random_field, random_mask, reshape, @@ -48,7 +49,7 @@ def _ccw_numpy( dy1dx2 = dy1 * dx2 lccw = np.where(dx1dy2 > dy1dx2, True, False) - ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise + ccw_out = np.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise return ccw_out @@ -109,8 +110,14 @@ def _line_intersect_numpy( line2_p2_lon, line2_p2_lat, ): - 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) + d1 = line1_p2_lon - line1_p1_lon + d1 = np.where(d1 != 0.0, d1, line1_p2_lon) + + d2 = line2_p2_lon - line2_p1_lon + d2 = np.where(d2 != 0.0, d2, line2_p2_lon) + + m1 = (line1_p2_lat - line1_p1_lat) / d1 + m2 = (line2_p2_lat - line2_p1_lat) / d2 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) @@ -966,7 +973,7 @@ def reference( True, False, ) - famask_bool = np.where(famask_int == int32(1), True, False) + famask_bool = np.where(famask_int == 1, True, False) mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) ps1_x, ps1_y = _line_intersect_numpy( fl_line_p1_lon, @@ -1482,19 +1489,19 @@ def reference( def input_data(self, grid): famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) p_vn = random_field(grid, EdgeDim, KDim) - ptr_v3_lon = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lon = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) - ptr_v3_lat = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lat = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) - tangent_orientation_dsl = random_field(grid, EdgeDim) - dreg_patch0_1_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_1_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lat_dsl = random_field(grid, EdgeDim, KDim) + tangent_orientation_dsl = random_field(grid, EdgeDim, low=0.1, high=1.0) + dreg_patch0_1_lon_dsl = constant_field(grid, 1.0, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = constant_field(grid, 1.0, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = constant_field(grid, 2.0, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = constant_field(grid, 2.0, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = constant_field(grid, 3.0, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = constant_field(grid, 3.0, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = constant_field(grid, 4.0, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = constant_field(grid, 4.0, EdgeDim, KDim) dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py similarity index 91% rename from model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py index 51fb24ffb9..b517a76c3c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py @@ -27,7 +27,7 @@ class TestHorAdvStencil01(StencilTest): PROGRAM = hor_adv_stencil_01 - OUTPUTS = ("tracer_new",) + OUTPUTS = ("tracer_new_hor",) @staticmethod def reference( @@ -43,13 +43,13 @@ def reference( ) -> np.array: geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) geofac_div = np.expand_dims(geofac_div, axis=-1) - tracer_new = ( + tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh * np.sum(p_mflx_tracer_h[grid.connectivities[C2EDim]] * geofac_div, axis=1) ) / rhodz_new - return dict(tracer_new=tracer_new) + return dict(tracer_new_hor=tracer_new_hor) @pytest.fixture def input_data(self, grid): @@ -61,7 +61,7 @@ def input_data(self, grid): geofac_div = random_field(grid, CellDim, C2EDim) geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) p_dtime = np.float64(5.0) - tracer_new = zero_field(grid, CellDim, KDim) + tracer_new_hor = zero_field(grid, CellDim, KDim) return dict( p_mflx_tracer_h=p_mflx_tracer_h, deepatmo_divh=deepatmo_divh, @@ -70,5 +70,5 @@ def input_data(self, grid): rhodz_new=rhodz_new, geofac_div=geofac_div_new, p_dtime=p_dtime, - tracer_new=tracer_new, + tracer_new_hor=tracer_new_hor, ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 95% rename from model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 05ca5d197a..1fe3317303 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -69,6 +69,7 @@ def reference( 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, 0, 1) p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( @@ -96,6 +97,7 @@ def reference( return dict( p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, @@ -122,6 +124,7 @@ def input_data(self, grid): dual_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) p_distv_bary_1 = random_field(grid, EdgeDim, KDim) p_distv_bary_2 = random_field(grid, EdgeDim, KDim) @@ -139,12 +142,13 @@ def input_data(self, grid): primal_normal_cell_2=primal_normal_cell_2_new, dual_normal_cell_2=dual_normal_cell_2_new, p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, p_dthalf=p_dthalf, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py index f883693ab1..213712810a 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -99,7 +99,7 @@ def _compute_wgt_t_detjac( 1.0 + zeta_4, ) - famask_bool = np.where(famask_int == np.int32(1), True, False) + famask_bool = np.where(famask_int == 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) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py similarity index 98% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py index dbb64d6807..6ab98e75ae 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -50,7 +50,7 @@ def reference( 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 == 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 @@ -65,7 +65,7 @@ def reference( @pytest.fixture def input_data(self, grid): - nsub = int32(1) + nsub = 1 p_mass_flx_e = random_field(grid, EdgeDim, KDim) geofac_div = random_field(grid, CellDim, C2EDim) z_rhofluxdiv_c = random_field(grid, CellDim, KDim) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py similarity index 95% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py index 59e3c026fe..185429af78 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py @@ -50,19 +50,19 @@ def reference( p_out_e = ( np.where( - cell_rel_idx_dsl == int32(1), + cell_rel_idx_dsl == 1, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0], ) + distv_bary_1 * np.where( - cell_rel_idx_dsl == int32(1), + cell_rel_idx_dsl == 1, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0], ) + distv_bary_2 * np.where( - cell_rel_idx_dsl == int32(1), + cell_rel_idx_dsl == 1, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0], ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py similarity index 69% rename from model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py index 598d496e3c..2e8c75481e 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py @@ -13,6 +13,8 @@ import numpy as np import pytest +from gt4py.next import as_field +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.model.common.dimension import CellDim, KDim @@ -32,14 +34,29 @@ def reference( deepatmo_divzl: np.array, deepatmo_divzu: np.array, rhodz_new: np.array, - p_dtime, + k: np.array, + ivadv_tracer: int32, + iadv_slev_jt: int32, + p_dtime: float, **kwargs, ) -> np.array: - tracer_new = ( - tracer_now * rhodz_now - + p_dtime - * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) - ) / rhodz_new + if ivadv_tracer != 0: + tracer_new = np.where( + (iadv_slev_jt <= k), + ( + tracer_now * rhodz_now + + p_dtime + * ( + p_mflx_tracer_v[:, 1:] * deepatmo_divzl + - p_mflx_tracer_v[:, :-1] * deepatmo_divzu + ) + ) + / rhodz_new, + tracer_now, + ) + else: + tracer_new = tracer_now + return dict(tracer_new=tracer_new) @pytest.fixture @@ -50,7 +67,10 @@ def input_data(self, grid): deepatmo_divzl = random_field(grid, KDim) deepatmo_divzu = random_field(grid, KDim) rhodz_new = random_field(grid, CellDim, KDim) + k = as_field((KDim,), np.arange(grid.num_levels, dtype=int32)) p_dtime = np.float64(5.0) + ivadv_tracer = 1 + iadv_slev_jt = 4 tracer_new = zero_field(grid, CellDim, KDim) return dict( tracer_now=tracer_now, @@ -59,6 +79,9 @@ def input_data(self, grid): deepatmo_divzl=deepatmo_divzl, deepatmo_divzu=deepatmo_divzu, rhodz_new=rhodz_new, + k=k, p_dtime=p_dtime, + ivadv_tracer=ivadv_tracer, + iadv_slev_jt=iadv_slev_jt, tracer_new=tracer_new, ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py similarity index 95% rename from model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py index f6fc3c2ff8..3474857e8d 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py @@ -28,7 +28,7 @@ class TestVLimitPrblSmStencil01(StencilTest): def reference(grid, p_face: np.array, p_cc: np.array, **kwargs): 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, 1, 0) return dict(l_limit=l_limit) @pytest.fixture diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py similarity index 98% rename from model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py index c79595b338..c85719d4ca 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py @@ -32,7 +32,7 @@ class TestVLimitPrblSmStencil02(StencilTest): @staticmethod def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs): q_face_up, q_face_low = np.where( - l_limit != int32(0), + l_limit != 0, np.where( (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), (p_cc, p_cc), diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py deleted file mode 100644 index 6b1daabc13..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py +++ /dev/null @@ -1,290 +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 -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( - divide_flux_area_list_stencil_02, -) -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - random_mask, - reshape, -) - - -class TestDivideFluxAreaListStencil02(StencilTest): - PROGRAM = divide_flux_area_list_stencil_02 - OUTPUTS = ( - "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", - ) - - @staticmethod - def reference( - grid, - 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, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - 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 = reshape(bf_cc_patch1_lon, e2c.shape) - bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) - bf_cc_patch1_lat = reshape(bf_cc_patch1_lat, e2c.shape) - bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) - bf_cc_patch2_lon = reshape(bf_cc_patch2_lon, e2c.shape) - bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) - bf_cc_patch2_lat = reshape(bf_cc_patch2_lat, e2c.shape) - 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 dict( - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask=patch1_cell_idx_vmask, - patch1_cell_blk_vmask=patch1_cell_blk_vmask, - patch2_cell_idx_vmask=patch2_cell_idx_vmask, - patch2_cell_blk_vmask=patch2_cell_blk_vmask, - ) - - @pytest.fixture - def input_data(self, grid): - famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_vn = random_field(grid, EdgeDim, KDim) - bf_cc_patch1_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) - bf_cc_patch1_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) - bf_cc_patch2_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) - bf_cc_patch2_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) - butterfly_idx_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - dreg_patch1_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lat_vmask = random_field(grid, EdgeDim, KDim) - patch1_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch1_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - - return dict( - famask_int=famask_int, - p_vn=p_vn, - bf_cc_patch1_lon=bf_cc_patch1_lon_field, - bf_cc_patch1_lat=bf_cc_patch1_lat_field, - bf_cc_patch2_lon=bf_cc_patch2_lon_field, - bf_cc_patch2_lat=bf_cc_patch2_lat_field, - butterfly_idx_patch1_vnpos=butterfly_idx_patch1_vnpos, - butterfly_idx_patch1_vnneg=butterfly_idx_patch1_vnneg, - butterfly_blk_patch1_vnpos=butterfly_blk_patch1_vnpos, - butterfly_blk_patch1_vnneg=butterfly_blk_patch1_vnneg, - butterfly_idx_patch2_vnpos=butterfly_idx_patch2_vnpos, - butterfly_idx_patch2_vnneg=butterfly_idx_patch2_vnneg, - butterfly_blk_patch2_vnpos=butterfly_blk_patch2_vnpos, - butterfly_blk_patch2_vnneg=butterfly_blk_patch2_vnneg, - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask=patch1_cell_idx_vmask, - patch1_cell_blk_vmask=patch1_cell_blk_vmask, - patch2_cell_idx_vmask=patch2_cell_idx_vmask, - patch2_cell_blk_vmask=patch2_cell_blk_vmask, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py deleted file mode 100644 index acfa579ff3..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ /dev/null @@ -1,184 +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 -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( - hflux_ffsl_hybrid_stencil_01a, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - constant_field, - random_field, - zero_field, -) - - -class TestHfluxFfslHybridStencil01a(StencilTest): - PROGRAM = hflux_ffsl_hybrid_stencil_01a - OUTPUTS = ("p_out_e_hybrid_1a",) - - @staticmethod - def reference( - grid, - 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_sum0_1: np.array, - z_quad_vector_sum0_2: np.array, - z_quad_vector_sum0_3: np.array, - z_quad_vector_sum0_4: np.array, - z_quad_vector_sum0_5: np.array, - z_quad_vector_sum0_6: np.array, - z_quad_vector_sum0_7: np.array, - z_quad_vector_sum0_8: np.array, - z_quad_vector_sum0_9: np.array, - z_quad_vector_sum0_10: np.array, - patch0_cell_rel_idx_dsl: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - 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 dict(p_out_e_hybrid_1a=p_out_e_hybrid_1a) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum0_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_10 = random_field(grid, EdgeDim, KDim) - patch0_cell_rel_idx_dsl = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) - p_out_e_hybrid_1a = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1=z_lsq_coeff_1, - z_lsq_coeff_2=z_lsq_coeff_2, - z_lsq_coeff_3=z_lsq_coeff_3, - z_lsq_coeff_4=z_lsq_coeff_4, - z_lsq_coeff_5=z_lsq_coeff_5, - z_lsq_coeff_6=z_lsq_coeff_6, - z_lsq_coeff_7=z_lsq_coeff_7, - z_lsq_coeff_8=z_lsq_coeff_8, - z_lsq_coeff_9=z_lsq_coeff_9, - z_lsq_coeff_10=z_lsq_coeff_10, - z_quad_vector_sum0_1=z_quad_vector_sum0_1, - z_quad_vector_sum0_2=z_quad_vector_sum0_2, - z_quad_vector_sum0_3=z_quad_vector_sum0_3, - z_quad_vector_sum0_4=z_quad_vector_sum0_4, - z_quad_vector_sum0_5=z_quad_vector_sum0_5, - z_quad_vector_sum0_6=z_quad_vector_sum0_6, - z_quad_vector_sum0_7=z_quad_vector_sum0_7, - z_quad_vector_sum0_8=z_quad_vector_sum0_8, - z_quad_vector_sum0_9=z_quad_vector_sum0_9, - z_quad_vector_sum0_10=z_quad_vector_sum0_10, - patch0_cell_rel_idx_dsl=patch0_cell_rel_idx_dsl, - p_out_e_hybrid_1a=p_out_e_hybrid_1a, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py deleted file mode 100644 index 0dbf30ca22..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.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 -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( - hflx_limiter_mo_stencil_01a, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil01a(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_01a - OUTPUTS = ( - "z_mflx_low", - "z_anti", - ) - - @staticmethod - def reference( - grid, p_mflx_tracer_h: np.ndarray, p_mass_flx_e: np.ndarray, p_cc: np.ndarray, **kwargs - ): - e2c = grid.connectivities[E2CDim] - 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 dict(z_mflx_low=z_mflx_low, z_anti=z_anti) - - @pytest.fixture - def input_data(self, grid): - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - z_mflx_low = zero_field(grid, EdgeDim, KDim) - z_anti = zero_field(grid, EdgeDim, KDim) - - return dict( - p_mflx_tracer_h=p_mflx_tracer_h, - p_mass_flx_e=p_mass_flx_e, - p_cc=p_cc, - z_mflx_low=z_mflx_low, - z_anti=z_anti, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py deleted file mode 100644 index f8819ee83e..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.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 -import pytest - -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.model.common.dimension import C2E2CDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil03MinMax(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_03_min_max - OUTPUTS = ("z_max", "z_min") - - @staticmethod - def reference(grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, **kwargs): - c2e2c = grid.connectivities[C2E2CDim] - 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 dict(z_max=z_max, z_min=z_min) - - @pytest.fixture - def input_data(self, grid): - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - beta_fct = 0.9 - r_beta_fct = 0.3 - z_max = zero_field(grid, CellDim, KDim) - z_min = zero_field(grid, CellDim, KDim) - return dict( - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - beta_fct=beta_fct, - r_beta_fct=r_beta_fct, - z_max=z_max, - z_min=z_min, - ) - - -class TestHflxLimiterMoStencil03(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_03 - OUTPUTS = ("r_p", "r_m") - - @staticmethod - def reference( - grid, - z_tracer_max, - z_tracer_min, - beta_fct, - r_beta_fct, - z_mflx_anti_in, - z_mflx_anti_out, - z_tracer_new_low, - dbl_eps, - **kwargs, - ): - res = TestHflxLimiterMoStencil03MinMax.reference( - grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct - ) - z_max, z_min = res["z_max"], res["z_min"] - 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 dict(r_p=r_p, r_m=r_m) - - @pytest.fixture - def input_data(self, grid): - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - beta_fct = 0.4 - r_beta_fct = 0.6 - z_mflx_anti_in = random_field(grid, CellDim, KDim) - z_mflx_anti_out = random_field(grid, CellDim, KDim) - z_tracer_new_low = random_field(grid, CellDim, KDim) - dbl_eps = 1e-5 - r_p = zero_field(grid, CellDim, KDim) - r_m = zero_field(grid, CellDim, KDim) - return dict( - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - beta_fct=beta_fct, - r_beta_fct=r_beta_fct, - z_mflx_anti_in=z_mflx_anti_in, - z_mflx_anti_out=z_mflx_anti_out, - z_tracer_new_low=z_tracer_new_low, - dbl_eps=dbl_eps, - r_p=r_p, - r_m=r_m, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py deleted file mode 100644 index d1b5e3244d..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_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 -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil04(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_04 - OUTPUTS = ("p_mflx_tracer_h",) - - @staticmethod - def reference( - grid, z_anti: np.ndarray, r_m: np.ndarray, r_p: np.ndarray, z_mflx_low: np.ndarray, **kwargs - ): - r_frac = np.where( - z_anti >= 0, - np.minimum( - r_m[grid.connectivities[E2CDim][:, 0]], r_p[grid.connectivities[E2CDim][:, 1]] - ), - np.minimum( - r_m[grid.connectivities[E2CDim][:, 1]], r_p[grid.connectivities[E2CDim][:, 0]] - ), - ) - return dict(p_mflx_tracer_h=z_mflx_low + np.minimum(1.0, r_frac) * z_anti) - - @pytest.fixture - def input_data(self, grid): - z_anti = random_field(grid, EdgeDim, KDim, low=-2.0, high=2.0) - r_m = random_field(grid, CellDim, KDim) - r_p = random_field(grid, CellDim, KDim) - z_mflx_low = random_field(grid, EdgeDim, KDim) - p_mflx_tracer_h = zero_field(grid, EdgeDim, KDim) - return dict( - z_anti=z_anti, r_m=r_m, r_p=r_p, z_mflx_low=z_mflx_low, p_mflx_tracer_h=p_mflx_tracer_h - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py deleted file mode 100644 index 212f95664a..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.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 -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field - - -class TestHflxLimiterPdStencil02(StencilTest): - PROGRAM = hflx_limiter_pd_stencil_02 - OUTPUTS = ("p_mflx_tracer_h",) - - @staticmethod - def reference(grid, refin_ctrl, r_m, p_mflx_tracer_h, bound, **kwargs): - e2c = grid.connectivities[E2CDim] - r_m_e2c = r_m[e2c] - refin_ctrl_expanded = np.expand_dims(refin_ctrl, axis=-1) - p_mflx_tracer_h_out = np.where( - refin_ctrl_expanded != bound, - np.where( - p_mflx_tracer_h >= 0, - p_mflx_tracer_h * r_m_e2c[:, 0], - p_mflx_tracer_h * r_m_e2c[:, 1], - ), - p_mflx_tracer_h, - ) - return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) - - @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) - def input_data(self, request, grid): - bound = np.int32(7) - scenario, ctrl_value = request.param - - if scenario == "no_match": - refin_ctrl = constant_field(grid, ctrl_value, EdgeDim, dtype=np.int32) - elif scenario == "everywhere_match": - refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) - elif scenario == "partly_match": - refin_ctrl = constant_field(grid, 5, EdgeDim, dtype=np.int32) - refin_ctrl[2:6] = bound - - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - - return dict( - refin_ctrl=refin_ctrl, - r_m=r_m, - p_mflx_tracer_h=p_mflx_tracer_h_in, - bound=bound, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py b/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py deleted file mode 100644 index 08da338d49..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py +++ /dev/null @@ -1,32 +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 pytest - -from icon4py.model.atmosphere.advection.init_zero_c import init_zero_c -from icon4py.model.common.dimension import CellDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestInitZeroCellK(StencilTest): - PROGRAM = init_zero_c - OUTPUTS = ("field",) - - @staticmethod - def reference(grid, **kwargs): - return dict(field=zero_field(grid, CellDim).asnumpy()) - - @pytest.fixture - def input_data(self, grid): - field = random_field(grid, CellDim) - return dict(field=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py deleted file mode 100644 index 47882226d9..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.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 - -import numpy as np -import pytest - -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 StencilTest, random_field, zero_field - - -class TestRbfIntpEdgeStencil01(StencilTest): - PROGRAM = rbf_intp_edge_stencil_01 - OUTPUTS = ("p_vt_out",) - - @staticmethod - def reference(grid, p_vn_in: np.array, ptr_coeff: np.array, **kwargs) -> np.array: - ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) - p_vt_out = np.sum(p_vn_in[grid.connectivities[E2C2EDim]] * ptr_coeff, axis=1) - return dict(p_vt_out=p_vt_out) - - @pytest.fixture - def input_data(self, grid): - p_vn_in = random_field(grid, EdgeDim, KDim) - ptr_coeff = random_field(grid, EdgeDim, E2C2EDim) - p_vt_out = zero_field(grid, EdgeDim, KDim) - return dict(p_vn_in=p_vn_in, ptr_coeff=ptr_coeff, p_vt_out=p_vt_out) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py deleted file mode 100644 index 298b6e71f3..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py +++ /dev/null @@ -1,77 +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 -import pytest - -from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( - recon_lsq_cell_l_svd_stencil, -) -from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -class TestReconLsqCellLSvdStencil(StencilTest): - PROGRAM = recon_lsq_cell_l_svd_stencil - OUTPUTS = ( - "p_coeff_1_dsl", - "p_coeff_2_dsl", - "p_coeff_3_dsl", - ) - - @staticmethod - def reference( - grid, - p_cc: np.array, - lsq_pseudoinv_1: np.array, - lsq_pseudoinv_2: np.array, - **kwargs, - ): - c2e2c = grid.connectivities[C2E2CDim] - p_cc_e = np.expand_dims(p_cc, axis=1) - n_diff = p_cc[c2e2c] - p_cc_e - lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c.shape) - lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c.shape) - lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) - p_coeff_1_dsl = p_cc - p_coeff_2_dsl = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) - p_coeff_3_dsl = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) - return dict( - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - ) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - lsq_pseudoinv_1_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) - lsq_pseudoinv_2_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) - p_coeff_1_dsl = zero_field(grid, CellDim, KDim) - p_coeff_2_dsl = zero_field(grid, CellDim, KDim) - p_coeff_3_dsl = zero_field(grid, CellDim, KDim) - return dict( - p_cc=p_cc, - lsq_pseudoinv_1=lsq_pseudoinv_1_field, - lsq_pseudoinv_2=lsq_pseudoinv_2_field, - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py deleted file mode 100644 index a538479e1e..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py +++ /dev/null @@ -1,194 +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 -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( - upwind_hflux_miura3_stencil_01, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - random_mask, - zero_field, -) - - -class TestUpwindHfluxMiura3Stencil01(StencilTest): - PROGRAM = upwind_hflux_miura3_stencil_01 - OUTPUTS = ("p_out_e_miura3",) - - @staticmethod - def reference( - grid, - 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, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - 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 dict(p_out_e_miura3=p_out_e_miura3) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_10 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - z_dreg_area = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_out_e_miura3 = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1=z_lsq_coeff_1, - z_lsq_coeff_2=z_lsq_coeff_2, - z_lsq_coeff_3=z_lsq_coeff_3, - z_lsq_coeff_4=z_lsq_coeff_4, - z_lsq_coeff_5=z_lsq_coeff_5, - z_lsq_coeff_6=z_lsq_coeff_6, - z_lsq_coeff_7=z_lsq_coeff_7, - z_lsq_coeff_8=z_lsq_coeff_8, - z_lsq_coeff_9=z_lsq_coeff_9, - z_lsq_coeff_10=z_lsq_coeff_10, - z_quad_vector_sum_1=z_quad_vector_sum_1, - z_quad_vector_sum_2=z_quad_vector_sum_2, - z_quad_vector_sum_3=z_quad_vector_sum_3, - z_quad_vector_sum_4=z_quad_vector_sum_4, - z_quad_vector_sum_5=z_quad_vector_sum_5, - z_quad_vector_sum_6=z_quad_vector_sum_6, - z_quad_vector_sum_7=z_quad_vector_sum_7, - z_quad_vector_sum_8=z_quad_vector_sum_8, - z_quad_vector_sum_9=z_quad_vector_sum_9, - z_quad_vector_sum_10=z_quad_vector_sum_10, - z_dreg_area=z_dreg_area, - p_mass_flx_e=p_mass_flx_e, - cell_rel_idx_dsl=cell_rel_idx_dsl, - p_out_e_miura3=p_out_e_miura3, - ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py deleted file mode 100644 index 1e33e6a7a7..0000000000 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ /dev/null @@ -1,92 +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 -import pytest -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.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - random_mask, - zero_field, -) - - -class TestUpwindHfluxMiuraCyclStencil01(StencilTest): - PROGRAM = upwind_hflux_miura_cycl_stencil_01 - OUTPUTS = ("z_tracer_mflx_dsl",) - - @staticmethod - def reference( - grid, - 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: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - 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 dict(z_tracer_mflx_dsl=z_tracer_mflx_dsl) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_2_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_3_dsl = random_field(grid, CellDim, KDim) - distv_bary_1 = random_field(grid, EdgeDim, KDim) - distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - z_tracer_mflx_dsl = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1_dsl=z_lsq_coeff_1_dsl, - z_lsq_coeff_2_dsl=z_lsq_coeff_2_dsl, - z_lsq_coeff_3_dsl=z_lsq_coeff_3_dsl, - distv_bary_1=distv_bary_1, - distv_bary_2=distv_bary_2, - p_mass_flx_e=p_mass_flx_e, - cell_rel_idx_dsl=cell_rel_idx_dsl, - z_tracer_mflx_dsl=z_tracer_mflx_dsl, - ) diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index cc70954c5e..bf85852f39 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -23,7 +23,8 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py-common>=0.0.5" + "icon4py-common>=0.0.5", + "nanobind<2.0.0" ] description = "ICON diffusion." dynamic = ['version'] diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 8dd62df576..3cfcb8945d 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -55,11 +55,11 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( k = broadcast(k, (CellDim, KDim)) dwdx, dwdy = ( where( - int32(0) < k, + 0 < k, _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) - if type_shear == int32(2) + if type_shear == 2 else (dwdx, dwdy) ) @@ -72,7 +72,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( ) w = where( - (int32(0) < k) & (k < nrdmax) & (interior_idx <= cell) & (cell < halo_idx), + (0 < k) & (k < nrdmax) & (interior_idx <= cell) & (cell < halo_idx), _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), w, ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index e2b8eb9f47..5709b4cfbf 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -52,9 +52,9 @@ def input_data(self, grid): z_nabla4_e2=z_nabla4_e2, diff_multfac_vn=diff_multfac_vn, vn=vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py index 2111f29011..7a3e550d2a 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -68,9 +68,9 @@ def input_data(self, grid): nudgecoeff_e=nudgecoeff_e, vn=vn, nudgezone_diff=nudgezone_diff, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index 8936a91542..684db54dcd 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -46,9 +46,9 @@ def input_data(self, grid): z_nabla2_e=z_nabla2_e, area_edge=area_edge, vn=vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py index 03f333d940..a1f2c0a9dc 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py @@ -67,8 +67,8 @@ def input_data(self, grid): geofac_n2s=geofac_n2s, w=w, diff_multfac_w=wpfloat("5.0"), - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 9b6852f436..4eec90119c 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -50,9 +50,9 @@ def input_data(self, grid): diff_multfac_n2w=diff_multfac_n2w, cell_area=cell_area, z_nabla2_c=z_nabla2_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py index 9449e08902..b0403b63a0 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -60,8 +60,8 @@ def input_data(self, grid): geofac_grg_y=geofac_grg_y, dwdx=dwdx, dwdy=dwdy, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py index defcdc80d4..cf04ca3305 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py @@ -50,8 +50,8 @@ def input_data(self, grid): w=w, geofac_n2s=geofac_n2s, z_nabla2_c=z_nabla2_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py index 2f062d616d..26001d74dc 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py @@ -71,8 +71,8 @@ def input_data(self, grid): inv_dual_edge_length=inv_dual_edge_length, theta_v=theta_v, z_nabla2_e=z_nabla2_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py index 0e37ca9b78..75e72ddc65 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py @@ -57,8 +57,8 @@ def input_data(self, grid): z_nabla2_e=z_nabla2_e, geofac_div=geofac_div_new, z_temp=z_temp, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py index 97ba066425..bdf348e658 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py @@ -126,8 +126,8 @@ def input_data(self, grid): inv_vert_vert_length=inv_vert_vert_length, inv_primal_edge_length=inv_primal_edge_length, z_nabla4_e2=z_nabla4_e2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 597638a112..130bd27eb8 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -45,8 +45,8 @@ def input_data(self, grid): return dict( kh_smag_e=kh_smag_e, enh_diffu_3d=enh_diffu_3d, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index dd80106e8e..a04c4b1583 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -64,8 +64,8 @@ def input_data(self, grid): enh_diffu_3d=enh_diffu_3d, thresh_tdiff=thresh_tdiff, smallest_vpfloat=smallest_vpfloat, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py index 5710ac13ad..e45018c211 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -73,8 +73,8 @@ def input_data(self, grid): diff_multfac_smag=diff_multfac_smag, kh_c=kh_c, div=div, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 2f9da07c1b..b44124de7a 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -139,8 +139,8 @@ def input_data(self, grid): theta_v=theta_v, z_temp=z_temp, vcoef=vcoef_new, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py index 7ed4bd1780..fda2e05221 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py @@ -69,8 +69,8 @@ def input_data(self, grid): theta_v=theta_v, exner=exner, rd_o_cvd=rd_o_cvd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index c51d649726..bab6f4cf93 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -23,7 +23,8 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py-common>=0.0.5" + "icon4py-common>=0.0.5", + "nanobind<2.0.0" ] description = "ICON dynamical core." dynamic = ['version'] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py index 76f0425ca8..7977851cef 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py @@ -60,12 +60,12 @@ def _fused_velocity_advection_stencil_16_to_18( k = broadcast(k, (CellDim, KDim)) ddt_w_adv = where( - (cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k), + (cell_lower_bound <= cell < cell_upper_bound) & (1 <= k), _compute_advective_vertical_wind_tendency(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz), ddt_w_adv, ) ddt_w_adv = where( - (cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k), + (cell_lower_bound <= cell < cell_upper_bound) & (1 <= k), _add_interpolated_horizontal_advection_of_w(e_bln_c_s, z_v_grad_w, ddt_w_adv), ddt_w_adv, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py index 9baf6d4e2f..58c9bc203c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py @@ -73,7 +73,7 @@ def compute_interface_vt_vn_and_kinetic_energy( ) (vn_ie, z_vt_ie, z_kin_hor_e) = where( - k == int32(0), + k == 0, _compute_horizontal_kinetic_energy(vn, vt), (vn_ie, z_vt_ie, z_kin_hor_e), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/init_exner_pr.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/init_exner_pr.py new file mode 100644 index 0000000000..5ec1cb398c --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/init_exner_pr.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 gt4py.next.common import GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32 + +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.settings import backend +from icon4py.model.common.type_alias import vpfloat + + +@field_operator +def _init_exner_pr( + exner: Field[[CellDim, KDim], vpfloat], + exner_ref: Field[[CellDim, KDim], vpfloat], +) -> Field[[CellDim, KDim], vpfloat]: + exner_pr = exner - exner_ref + return exner_pr + + +@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +def init_exner_pr( + exner: Field[[CellDim, KDim], vpfloat], + exner_ref: Field[[CellDim, KDim], vpfloat], + exner_pr: Field[[CellDim, KDim], vpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _init_exner_pr( + exner, + exner_ref, + out=exner_pr, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index c7ce4a55c2..ae6df8f9ac 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -1647,20 +1647,20 @@ def run_corrector_step( offset_provider={}, ) - # verified for e-9 - log.debug(f"corrector: start stencile 41") - compute_divergence_of_fluxes_of_rho_and_theta( - geofac_div=self.interpolation_state.geofac_div, - mass_fl_e=diagnostic_state_nh.mass_fl_e, - z_theta_v_fl_e=self.z_theta_v_fl_e, - z_flxdiv_mass=self.z_flxdiv_mass, - z_flxdiv_theta=self.z_flxdiv_theta, - horizontal_start=start_cell_nudging, - horizontal_end=end_cell_local, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider=self.grid.offset_providers, - ) + # verified for e-9 + log.debug(f"corrector: start stencil 41") + compute_divergence_of_fluxes_of_rho_and_theta( + geofac_div=self.interpolation_state.geofac_div, + mass_fl_e=diagnostic_state_nh.mass_fl_e, + z_theta_v_fl_e=self.z_theta_v_fl_e, + z_flxdiv_mass=self.z_flxdiv_mass, + z_flxdiv_theta=self.z_flxdiv_theta, + horizontal_start=start_cell_nudging, + horizontal_end=end_cell_local, + vertical_start=0, + vertical_end=self.grid.num_levels, + offset_provider=self.grid.offset_providers, + ) if self.config.itime_scheme == 4: log.debug(f"corrector start stencil 42 44 45 45b") diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py index 2fb74f7f36..643d71c47e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py @@ -119,7 +119,7 @@ def _predictor_stencils_2_3( nlev: int32, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: (z_exner_ex_pr, exner_pr) = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _extrapolate_temporally_exner_pressure(exner_exfac, exner, exner_ref_mc, exner_pr), (z_exner_ex_pr, exner_pr), ) @@ -253,13 +253,13 @@ def _predictor_stencils_7_8_9( Field[[CellDim, KDim], float], ]: (z_rth_pr_1, z_rth_pr_2) = where( - k_field == int32(0), + k_field == 0, _compute_perturbation_of_rho_and_theta(rho, rho_ref_mc, theta_v, theta_ref_mc), (z_rth_pr_1, z_rth_pr_2), ) (rho_ic, z_rth_pr_1, z_rth_pr_2) = where( - k_field >= int32(1), + k_field >= 1, _compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers( wgtfac_c, rho, rho_ref_mc, theta_v, theta_ref_mc ), @@ -267,7 +267,7 @@ def _predictor_stencils_7_8_9( ) (z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c) = where( - k_field >= int32(1), + k_field >= 1, _compute_virtual_potential_temperatures_and_pressure_gradient( wgtfac_c, z_rth_pr_2, @@ -350,9 +350,7 @@ def _predictor_stencils_11_lower_upper( k_field: Field[[KDim], int32], nlev: int32, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_theta_v_pr_ic = where( - k_field == int32(0), _init_cell_kdim_field_with_zero_vp(), z_theta_v_pr_ic - ) + z_theta_v_pr_ic = where(k_field == 0, _init_cell_kdim_field_with_zero_vp(), z_theta_v_pr_ic) (z_theta_v_pr_ic, theta_v_ic) = where( k_field == nlev, @@ -469,7 +467,7 @@ def _predictor_stencils_35_36( z_w_concorr_me, ) (vn_ie, z_vt_ie, z_kin_hor_e) = where( - k_field >= int32(1), + k_field >= 1, _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges(wgtfac_e, vn, vt), (vn_ie, z_vt_ie, z_kin_hor_e), ) @@ -643,7 +641,7 @@ def _stencils_42_44_45_45b( Field[[CellDim, KDim], float], ]: (z_w_expl, z_contr_w_fl_l) = where( - (k_field >= int32(1)) & (k_field < nlev), + (k_field >= 1) & (k_field < nlev), _compute_explicit_vertical_wind_from_advection_and_vertical_wind_density( w_nnow, ddt_w_adv_ntl1, @@ -661,7 +659,7 @@ def _stencils_42_44_45_45b( ) (z_beta, z_alpha) = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _compute_solver_coefficients_matrix( exner_nnow, rho_nnow, @@ -678,7 +676,7 @@ def _stencils_42_44_45_45b( ) z_alpha = where(k_field == nlev, _init_cell_kdim_field_with_zero_vp(), z_alpha) - z_q = where(k_field == int32(0), _init_cell_kdim_field_with_zero_vp(), z_q) + z_q = where(k_field == 0, _init_cell_kdim_field_with_zero_vp(), z_q) return z_w_expl, z_contr_w_fl_l, z_beta, z_alpha, z_q @@ -783,7 +781,7 @@ def _stencils_43_44_45_45b( Field[[CellDim, KDim], float], ]: (z_w_expl, z_contr_w_fl_l) = where( - (k_field >= int32(1)) & (k_field < nlev), + (k_field >= 1) & (k_field < nlev), _compute_explicit_vertical_wind_speed_and_vertical_wind_times_density( w_nnow, ddt_w_adv_ntl1, @@ -797,7 +795,7 @@ def _stencils_43_44_45_45b( (z_w_expl, z_contr_w_fl_l), ) (z_beta, z_alpha) = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _compute_solver_coefficients_matrix( exner_nnow, rho_nnow, @@ -813,7 +811,7 @@ def _stencils_43_44_45_45b( (z_beta, z_alpha), ) z_alpha = where(k_field == nlev, _init_cell_kdim_field_with_zero_vp(), z_alpha) - z_q = where(k_field == int32(0), _init_cell_kdim_field_with_zero_vp(), z_q) + z_q = where(k_field == 0, _init_cell_kdim_field_with_zero_vp(), z_q) return z_w_expl, z_contr_w_fl_l, z_beta, z_alpha, z_q @@ -911,7 +909,7 @@ def _stencils_47_48_49( ) # 48 and 49 are identical except for bounds (z_rho_expl, z_exner_expl) = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _compute_explicit_part_for_rho_and_exner( rho_nnow, inv_ddqz_z_full, @@ -997,7 +995,7 @@ def _stencils_61_62( Field[[CellDim, KDim], float], ]: (rho_new, exner_new, w_new) = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _update_density_exner_wind( rho_now, grf_tend_rho, theta_v_now, grf_tend_thv, w_now, grf_tend_w, dtime ), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py index a3751eeb64..5daef8762a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py @@ -71,7 +71,7 @@ def _fused_stencils_4_5( ) (vn_ie, z_vt_ie, z_kin_hor_e) = where( - k_field == int32(0), + k_field == 0, _compute_horizontal_kinetic_energy(vn, vt), (vn_ie, z_vt_ie, z_kin_hor_e), ) @@ -156,7 +156,7 @@ def _fused_stencils_9_10( ) w_concorr_c = where( - (k_field >= nflatlev_startindex + int32(1)) & (k_field < nlev), + (k_field >= nflatlev_startindex + 1) & (k_field < nlev), _interpolate_to_half_levels_vp(interpolant=local_z_w_concorr_mc, wgtfac_c=wgtfac_c), w_concorr_c, ) @@ -206,7 +206,7 @@ def _fused_stencils_11_to_13( nlev: int32, ): local_z_w_con_c = where( - (k_field >= int32(0)) & (k_field < nlev), + (k_field >= 0) & (k_field < nlev), _copy_cell_kdim_field_to_vp(w), local_z_w_con_c, ) @@ -214,7 +214,7 @@ def _fused_stencils_11_to_13( local_z_w_con_c = where(k_field == nlev, _init_cell_kdim_field_with_zero_vp(), local_z_w_con_c) local_z_w_con_c = where( - (k_field >= (nflatlev_startindex + int32(1))) & (k_field < nlev), + (k_field >= (nflatlev_startindex + 1)) & (k_field < nlev), _correct_contravariant_vertical_velocity(local_z_w_con_c, w_concorr_c), local_z_w_con_c, ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_accumulate_prep_adv_fields.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_accumulate_prep_adv_fields.py index fd62b0edb5..cf7a4d3078 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_accumulate_prep_adv_fields.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_accumulate_prep_adv_fields.py @@ -53,8 +53,8 @@ def input_data(self, grid): vn_traj=vn_traj, mass_flx_me=mass_flx_me, r_nsubsteps=r_nsubsteps, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py index 96066bebb8..6c59ad9bd2 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py @@ -55,8 +55,8 @@ def input_data(self, grid): rho_incr=rho_incr, exner_incr=exner_incr, iau_wgt_dyn=iau_wgt_dyn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_to_vn.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_to_vn.py index 533ac55a7a..82af8ca3be 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_to_vn.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_to_vn.py @@ -42,8 +42,8 @@ def input_data(self, grid): vn_incr=vn_incr, vn=vn, iau_wgt_dyn=iau_wgt_dyn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py index e6266ec0ec..9764c286ab 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py @@ -141,9 +141,9 @@ def input_data(self, grid): dtime=dtime, c_lin_e=c_lin_e, z_w_con_c_full=z_w_con_c_full, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py index ab1dfe535c..dab1718e28 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py @@ -137,8 +137,8 @@ def input_data(self, grid): scalfac_exdiff=scalfac_exdiff, cfl_w_limit=cfl_w_limit, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py index a8ae9fc9b2..ce8cf650ba 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py @@ -57,8 +57,8 @@ def input_data(self, grid): e_bln_c_s=e_bln_c_s, z_v_grad_w=z_v_grad_w, ddt_w_adv=ddt_w_adv, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn.py index 5e80f37361..c3a4ad816a 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn.py @@ -63,8 +63,8 @@ def input_data(self, grid): vn_nnew=vn_nnew, dtime=dtime, cpd=cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py index de9359d8fc..615ec63c18 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py @@ -76,8 +76,8 @@ def input_data(self, grid): wgt_nnow_vel=wgt_nnow_vel, wgt_nnew_vel=wgt_nnew_vel, cpd=cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py index 1448236f0d..95b73c022d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py @@ -66,8 +66,8 @@ def input_data(self, grid): inv_dual_edge_length=inv_dual_edge_length, z_dwdz_dd=z_dwdz_dd, z_graddiv_vn=z_graddiv_vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_2nd_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_2nd_order_divergence_damping.py index 3e10b25cd5..3917f3c009 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_2nd_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_2nd_order_divergence_damping.py @@ -42,8 +42,8 @@ def input_data(self, grid): z_graddiv_vn=z_graddiv_vn, vn=vn, scal_divdamp_o2=scal_divdamp_o2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py index 72caaed92c..d82a9c4262 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py @@ -49,8 +49,8 @@ def input_data(self, grid): scal_divdamp=scal_divdamp, z_graddiv2_vn=z_graddiv2_vn, vn=vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py index 26e48e1317..c7c16a87f2 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py @@ -54,8 +54,8 @@ def input_data(self, grid): pg_exdist=pg_exdist, z_hydro_corr=z_hydro_corr, z_gradh_exner=z_gradh_exner, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py index fb0cf09c08..477cd752b7 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py @@ -44,8 +44,8 @@ def input_data(self, grid): z_raylfac=z_raylfac, w_1=w_1, w=w, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py index 2e66c0cb9e..885ead1075 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -55,8 +55,8 @@ def input_data(self, grid): nudgecoeff_e=nudgecoeff_e, z_graddiv2_vn=z_graddiv2_vn, vn=vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_normal_wind_tendency.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_normal_wind_tendency.py index 3b21e8bdd5..fe4f73db82 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_normal_wind_tendency.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_normal_wind_tendency.py @@ -123,8 +123,8 @@ def input_data(self, grid): z_w_con_c_full=z_w_con_c_full, ddqz_z_full_e=ddqz_z_full_e, ddt_vn_apc=ddt_vn_apc, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_vertical_wind_tendency.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_vertical_wind_tendency.py index 6c93461d84..f7bfeec8a7 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_vertical_wind_tendency.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_vertical_wind_tendency.py @@ -71,8 +71,8 @@ def input_data(self, grid): coeff1_dwdz=coeff1_dwdz, coeff2_dwdz=coeff2_dwdz, ddt_w_adv=ddt_w_adv, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py index fb7b666269..0fbf5694dd 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py @@ -43,8 +43,8 @@ def input_data(self, grid): ddqz_z_full_in=ddqz_z_full_in, deepatmo_t1mc_in=deepatmo_t1mc_in, airmass_out=airmass_out, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py index 83e0e14ec2..f7ed5b5917 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py @@ -58,8 +58,8 @@ def input_data(self, grid): d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, z_rth_pr_2=z_rth_pr_2, z_dexner_dz_c_2=z_dexner_dz_c_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn.py index 03645a5d6f..e00ade97e6 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn.py @@ -44,8 +44,8 @@ def input_data(self, grid): e_flx_avg=e_flx_avg, vn=vn, z_vn_avg=z_vn_avg, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py index 3887347fc3..7a26e94acf 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py @@ -70,8 +70,8 @@ def input_data(self, grid): z_vn_avg=z_vn_avg, z_graddiv_vn=z_graddiv_vn, vt=vt, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py index d6f29df4ba..95f2c3773f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py @@ -60,8 +60,8 @@ def input_data(self, grid): ddxt_z_full=ddxt_z_full, vt=vt, z_w_concorr_me=z_w_concorr_me, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w.py index c69993b56a..444340aad2 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w.py @@ -68,8 +68,8 @@ def input_data(self, grid): z_w_concorr_me=z_w_concorr_me, wgtfac_c=wgtfac_c, w_concorr_c=w_concorr_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py index d3a6953221..b35f56198b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py @@ -78,7 +78,7 @@ def input_data(self, grid): z_w_concorr_me=z_w_concorr_me, wgtfacq_c=wgtfacq_c, w_concorr_c=w_concorr_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), vertical_start=int32(grid.num_levels - 1), vertical_end=int32(grid.num_levels), diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py index 33a002569e..75363ed799 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py @@ -48,8 +48,8 @@ def input_data(self, grid): w=w, w_concorr_c=w_concorr_c, z_dwdz_dd=z_dwdz_dd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py index 4a106dd1fe..5154f214f8 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py @@ -52,8 +52,8 @@ def input_data(self, grid): exner=exner, rd_o_cvd=rd_o_cvd, rd_o_p0ref=rd_o_p0ref, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_part_for_rho_and_exner.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_part_for_rho_and_exner.py index 54d3972ca0..51f4a17d64 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_part_for_rho_and_exner.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_part_for_rho_and_exner.py @@ -87,8 +87,8 @@ def input_data(self, grid): theta_v_ic=theta_v_ic, ddt_exner_phy=ddt_exner_phy, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py index 7069073fbc..f3faa2856a 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py @@ -80,8 +80,8 @@ def input_data(self, grid): wgt_nnow_vel=wgt_nnow_vel, wgt_nnew_vel=wgt_nnew_vel, cpd=cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py index 543669010b..d368c3f3b4 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py @@ -69,8 +69,8 @@ def input_data(self, grid): vwind_expl_wgt=vwind_expl_wgt, dtime=dtime, cpd=cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_first_vertical_derivative.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_first_vertical_derivative.py index 5fa288801c..00d0ac359b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_first_vertical_derivative.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_first_vertical_derivative.py @@ -42,8 +42,8 @@ def input_data(self, grid): z_exner_ic=z_exner_ic, inv_ddqz_z_full=inv_ddqz_z_full, z_dexner_dz_c_1=z_dexner_dz_c_1, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_graddiv2_of_vn.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_graddiv2_of_vn.py index 2a4e9c0303..7cc7af203b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_graddiv2_of_vn.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_graddiv2_of_vn.py @@ -45,8 +45,8 @@ def input_data(self, grid): geofac_grdiv=geofac_grdiv, z_graddiv_vn=z_graddiv_vn, z_graddiv2_vn=z_graddiv2_vn, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py index 1f34f5202d..c45b89182b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py @@ -237,8 +237,8 @@ def input_data(self, grid): z_rth_pr_2=z_rth_pr_2, z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py index 9323095aa0..c768906d76 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py @@ -101,8 +101,8 @@ def input_data(self, grid): tangent_orientation=tangent_orientation, z_w_v=z_w_v, z_v_grad_w=z_v_grad_w, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py index d994f9b904..27b6140e4d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py @@ -50,8 +50,8 @@ def input_data(self, grid): inv_dual_edge_length=inv_dual_edge_length, z_exner_ex_pr=z_exner_ex_pr, z_gradh_exner=z_gradh_exner, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py index 0186ec37df..351e9ee252 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py @@ -109,8 +109,8 @@ def input_data(self, grid): z_dexner_dz_c_1=z_dexner_dz_c_1, z_dexner_dz_c_2=z_dexner_dz_c_2, z_gradh_exner=z_gradh_exner, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py index a487242229..664cdbd28b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py @@ -68,8 +68,8 @@ def input_data(self, grid): c_lin_e=c_lin_e, z_dexner_dz_c_1=z_dexner_dz_c_1, z_gradh_exner=z_gradh_exner, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_kinetic_energy.py index 59e04a8690..747a22f24c 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_kinetic_energy.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_horizontal_kinetic_energy.py @@ -54,8 +54,8 @@ def input_data(self, grid): vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_hydrostatic_correction_term.py index a2dcd267bb..76ec24434f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_hydrostatic_correction_term.py @@ -134,8 +134,8 @@ def input_data(self, grid): inv_ddqz_z_full=inv_ddqz_z_full, inv_dual_edge_length=inv_dual_edge_length, grav_o_cpd=grav_o_cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_mass_flux.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_mass_flux.py index 0f9fb7fbbb..76ad0ccb09 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_mass_flux.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_mass_flux.py @@ -54,8 +54,8 @@ def input_data(self, grid): z_theta_v_e=z_theta_v_e, mass_fl_e=mass_fl_e, z_theta_v_fl_e=z_theta_v_fl_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py index 940b8bf5d6..0a9c36ed8b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py @@ -89,8 +89,8 @@ def input_data(self, grid): vcfl=vcfl, cfl_w_limit=cfl_w_limit, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta.py index c82bedd3fd..6b76c4c2ea 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta.py @@ -56,8 +56,8 @@ def input_data(self, grid): theta_ref_mc=theta_ref_mc, z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py index 0fd1de6220..81dd425938 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py @@ -47,9 +47,9 @@ def input_data(self, grid): rho_ic=rho_ic, z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py index e3c97e6dcd..f6bc7a8472 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py @@ -102,8 +102,8 @@ def input_data(self, grid): theta_v_new=theta_v_new, dtime=dtime, cvd_o_rd=cvd_o_rd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py index 1695400895..7211ded64e 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py @@ -129,8 +129,8 @@ def input_data(self, grid): dtime=dtime, wgt_nnow_rth=wgt_nnow_rth, wgt_nnew_rth=wgt_nnew_rth, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_solver_coefficients_matrix.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_solver_coefficients_matrix.py index f303d1a4b3..3d06950a98 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_solver_coefficients_matrix.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_solver_coefficients_matrix.py @@ -76,8 +76,8 @@ def input_data(self, grid): dtime=dtime, rd=rd, cvd=cvd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_tangential_wind.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_tangential_wind.py index 1fc2e5768f..1afd46c694 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_tangential_wind.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_tangential_wind.py @@ -47,8 +47,8 @@ def input_data(self, grid): vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e, vt=vt, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py index bfd08ca2d3..0eb3209e8d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py @@ -61,8 +61,8 @@ def input_data(self, grid): exner=exner, rd_o_cvd=rd_o_cvd, rd_o_p0ref=rd_o_p0ref, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py index 6e530ecad6..a6db2593da 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -83,8 +83,8 @@ def input_data(self, grid): z_theta_v_pr_ic=z_theta_v_pr_ic, theta_v_ic=theta_v_ic, z_th_ddz_exner_c=z_th_ddz_exner_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_vn_on_lateral_boundary.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_vn_on_lateral_boundary.py index db7df4e663..1f5af20b91 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_vn_on_lateral_boundary.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_vn_on_lateral_boundary.py @@ -44,8 +44,8 @@ def input_data(self, grid): vn_now=vn_now, vn_new=vn_new, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_copy_cell_kdim_field_to_vp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_copy_cell_kdim_field_to_vp.py index fb7cdccbe5..a17269a466 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_copy_cell_kdim_field_to_vp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_copy_cell_kdim_field_to_vp.py @@ -42,8 +42,8 @@ def input_data(self, grid): return dict( field=field, field_copy=field_copy, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_correct_contravariant_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_correct_contravariant_vertical_velocity.py index 208af69657..c3b2c4590c 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_correct_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_correct_contravariant_vertical_velocity.py @@ -47,8 +47,8 @@ def input_data(self, grid): return dict( w_concorr_c=w_concorr_c, z_w_con_c=z_w_con_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py index 7a3b7465fb..1dea12b413 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py @@ -56,7 +56,7 @@ def input_data(self, grid): wgtfacq_e=wgtfacq_e, vn=vn, vn_ie=vn_ie, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), vertical_start=int32(grid.num_levels), vertical_end=int32(grid.num_levels + 1), diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_temporally_exner_pressure.py index 60228271c0..70f5d12c35 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_temporally_exner_pressure.py @@ -54,8 +54,8 @@ def input_data(self, grid): exner_ref_mc=exner_ref_mc, exner_pr=exner_pr, z_exner_ex_pr=z_exner_ex_pr, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py index 39c1c9b95e..4fcd423744 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py @@ -89,7 +89,7 @@ def input_data(self, grid): nlev=nlev, nflatlev=nflatlev, w_concorr_c=w_concorr_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), vertical_start=int32(grid.num_levels - 1), vertical_end=int32(grid.num_levels), diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_vp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_vp.py index 94bf7156c2..f8661f6323 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_vp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_vp.py @@ -45,8 +45,8 @@ def input_data(self, grid): return dict( field_with_zero_vp=field_with_zero_vp, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index 1df051631c..c1fa6da96c 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -38,8 +38,8 @@ def input_data(self, grid): return dict( field_with_zero_wp=field_with_zero_wp, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_exner_pr.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_exner_pr.py new file mode 100644 index 0000000000..a40dd7e5bd --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_exner_pr.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 gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.dycore.init_exner_pr import ( + init_exner_pr, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.type_alias import vpfloat + + +class TestInitExnerPr(StencilTest): + PROGRAM = init_exner_pr + OUTPUTS = ("exner_pr",) + + @staticmethod + def reference(grid, exner: np.array, exner_ref: np.array, **kwargs) -> dict: + exner_pr = exner - exner_ref + return dict( + exner_pr=exner_pr, + ) + + @pytest.fixture + def input_data(self, grid): + exner = random_field(grid, CellDim, KDim, dtype=vpfloat) + exner_ref = random_field(grid, CellDim, KDim, dtype=vpfloat) + exner_pr = zero_field(grid, CellDim, KDim, dtype=vpfloat) + + return dict( + exner=exner, + exner_ref=exner_ref, + exner_pr=exner_pr, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_cells), + vertical_start=int32(0), + vertical_end=int32(grid.num_levels), + ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_index_with_zero_vp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_index_with_zero_vp.py index fb432f1f96..280b22b26f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_index_with_zero_vp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_index_with_zero_vp.py @@ -55,7 +55,7 @@ def input_data(self, grid): field_index_with_zero_2 = random_field(grid, CellDim, KDim, dtype=vpfloat) k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) - k1 = int32(1) + k1 = 1 k2 = int32(grid.num_levels) return dict( @@ -64,8 +64,8 @@ def input_data(self, grid): k=k, k1=k1, k2=k2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py index 0fc8327710..5aedc61442 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py @@ -49,8 +49,8 @@ def input_data(self, grid): return dict( cell_kdim_field_with_zero_vp_1=cell_kdim_field_with_zero_vp_1, cell_kdim_field_with_zero_vp_2=cell_kdim_field_with_zero_vp_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_wp.py index e9e0abe79f..61299a188e 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_cell_kdim_fields_with_zero_wp.py @@ -49,8 +49,8 @@ def input_data(self, grid): return dict( cell_kdim_field_with_zero_wp_1=cell_kdim_field_with_zero_wp_1, cell_kdim_field_with_zero_wp_2=cell_kdim_field_with_zero_wp_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py index 571d5ec5d9..83c95aaf56 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py @@ -49,8 +49,8 @@ def input_data(self, grid): return dict( edge_kdim_field_with_zero_wp_1=edge_kdim_field_with_zero_wp_1, edge_kdim_field_with_zero_wp_2=edge_kdim_field_with_zero_wp_2, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py index 957258237e..864947a258 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py @@ -48,8 +48,8 @@ def input_data(self, grid): return dict( z_w_con_c=z_w_con_c, z_w_con_c_full=z_w_con_c_full, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py index f9778c45c6..19ce7bbac2 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py @@ -57,8 +57,8 @@ def input_data(self, grid): interpolant=interpolant, e_bln_c_s=as_1D_sparse_field(e_bln_c_s, CEDim), interpolation=interpolation, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_half_levels_vp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_half_levels_vp.py index 62d61effe0..830641a691 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_half_levels_vp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_half_levels_vp.py @@ -56,8 +56,8 @@ def input_data(self, grid): wgtfac_c=wgtfac_c, interpolant=interpolant, interpolation_to_half_levels_vp=interpolation_to_half_levels_vp, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_surface.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_surface.py index 6239080a5b..31daa0a702 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_surface.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_surface.py @@ -62,8 +62,8 @@ def input_data(self, grid): interpolant=interpolant, wgtfacq_c=wgtfacq_c, interpolation_to_surface=interpolation_to_surface, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(3), + vertical_start=3, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py index 6edf7ffb88..2f8d75b2f9 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py @@ -95,8 +95,8 @@ def input_data(self, grid): vt=vt, vn_ie=vn_ie, z_kin_hor_e=z_kin_hor_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_interface_edges.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_interface_edges.py index ffa3b691c7..ac7b9216da 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_interface_edges.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_interface_edges.py @@ -63,8 +63,8 @@ def input_data(self, grid): wgtfac_e=wgtfac_e, vt=vt, z_vt_ie=z_vt_ie, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mcompute_divergence_of_fluxes_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mcompute_divergence_of_fluxes_of_rho_and_theta.py index 5998e4b3c4..f1fff490a3 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mcompute_divergence_of_fluxes_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mcompute_divergence_of_fluxes_of_rho_and_theta.py @@ -66,8 +66,8 @@ def input_data(self, grid): z_theta_v_fl_e=z_theta_v_fl_e, z_flxdiv_mass=z_flxdiv_mass, z_flxdiv_theta=z_flxdiv_theta, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index fb1dcfe6ba..766e23646d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -55,8 +55,8 @@ def input_data(self, grid): p_cell_in=p_cell_in, c_intp=c_intp, p_vert_out=p_vert_out, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_vertices), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py index d44025b13a..1b121b2c33 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py @@ -49,8 +49,8 @@ def input_data(self, grid): vec_e=vec_e, geofac_rot=geofac_rot, rot_vec=rot_vec, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_vertices), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py index 47c524b99a..35d21529fc 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py @@ -78,8 +78,8 @@ def input_data(self, grid): p_ccpr2=p_ccpr2, geofac_grg_x=geofac_grg_x, geofac_grg_y=geofac_grg_y, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py index 00b903e42c..852f893275 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py @@ -100,8 +100,8 @@ def input_data(self, grid): z_exner_expl=z_exner_expl, dtime=dtime, cpd=cpd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py index 3b5250bfa8..d32145a6af 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py @@ -43,8 +43,8 @@ def input_data(self, grid): w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l, w_concorr_c=w_concorr_c, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py index 34543761aa..91e8a0eb71 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py @@ -62,8 +62,8 @@ def input_data(self, grid): theta_ref_ic=theta_ref_ic, z_theta_v_pr_ic=z_theta_v_pr_ic, theta_v_ic=theta_v_ic, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(3), + vertical_start=3, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py index 04986df2f7..1444262b9d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py @@ -42,9 +42,9 @@ def reference(grid, z_q: np.array, w: np.array, **kwargs) -> dict: def input_data(self, grid): z_q = random_field(grid, CellDim, KDim, dtype=vpfloat) w = random_field(grid, CellDim, KDim, dtype=wpfloat) - h_start = int32(0) + h_start = 0 h_end = int32(grid.num_cells) - v_start = int32(1) + v_start = 1 v_end = int32(grid.num_levels) return dict( z_q=z_q, diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py index ebf5df38c0..38d47230f8 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -88,9 +88,9 @@ def test_solve_tridiagonal_matrix_for_w_forward_sweep(): dtime, cpd, ) - h_start = int32(0) + h_start = 0 h_end = int32(grid.num_cells) - v_start = int32(1) + v_start = 1 v_end = int32(grid.num_levels) # TODO we run this test with the C++ backend as the `embedded` backend doesn't handle this pattern solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(run_gtfn)( diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_density_exner_wind.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_density_exner_wind.py index 21b4e8767f..9f12f82972 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_density_exner_wind.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_density_exner_wind.py @@ -66,8 +66,8 @@ def input_data(self, grid): exner_new=exner_new, w_new=w_new, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py index a976063824..4bbc360dad 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py @@ -53,8 +53,8 @@ def input_data(self, grid): exner_dyn_incr=exner_dyn_incr, ndyn_substeps_var=ndyn_substeps_var, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py index c6b2bd6995..d2e842e7ad 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py @@ -65,8 +65,8 @@ def input_data(self, grid): w_concorr_c=w_concorr_c, mass_flx_ic=mass_flx_ic, r_nsubsteps=r_nsubsteps, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py index a414a05f3b..4d1abde01f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py @@ -64,8 +64,8 @@ def input_data(self, grid): mass_flx_ic=mass_flx_ic, vol_flx_ic=vol_flx_ic, r_nsubsteps=r_nsubsteps, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py index 8e861b2cf1..300ed2bc3f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py @@ -67,8 +67,8 @@ def input_data(self, grid): rho_new=rho_new, theta_v_new=theta_v_new, cvd_o_rd=cvd_o_rd, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py index 6644e77137..6d29b37ef8 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py @@ -42,8 +42,8 @@ def input_data(self, grid): grf_tend_w=grf_tend_w, w_new=w_new, dtime=dtime, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/testinterpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/testinterpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py index f03a0d0090..9591de9785 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/testinterpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/testinterpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py @@ -58,8 +58,8 @@ def input_data(self, grid): vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(grid.num_levels), ) diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 69ce8ba965..c9dfc455da 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -19,7 +19,8 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1" + "gt4py>=1.0.1", + "nanobind<2.0.0" ] description = "Shared code for the icon4py model." dynamic = ['version'] diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 88ca676dfe..484974f2c5 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -45,6 +45,8 @@ _ICON_INDEX_OFFSET_EDGES: Final[int] = 13 _GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 +_GRF_NUDGEZONE_START_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + 1 +_GRF_NUDGEZONE_WIDTH: Final[int] = 8 _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 diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 28106d97a7..5f1f69f37a 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -86,7 +86,7 @@ def _determine_kstart_moist( def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): assert damping_height >= 0.0, "Damping height must be positive." return ( - int32(0) + 0 if damping_height > vct_a[0] else int32(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) ) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 9c1fbb72e1..0f2a12265c 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -10,6 +10,8 @@ # 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 import ( Field, @@ -57,6 +59,12 @@ """ +@dataclass(frozen=True) +class MetricsConfig: + #: Temporal extrapolation of Exner for computation of horizontal pressure gradient, defined in `mo_nonhydrostatic_nml.f90` used only in metrics fields calculation. + exner_expol: Final[wpfloat] = 0.333 + + @program(grid_type=GridType.UNSTRUCTURED) def compute_z_mc( z_ifc: Field[[CellDim, KDim], wpfloat], @@ -361,7 +369,7 @@ def _compute_d2dexdz2_fac1_mc( grav: wpfloat, igradp_method: int32, ) -> Field[[CellDim, KDim], vpfloat]: - if igradp_method <= int32(3): + if igradp_method <= 3: d2dexdz2_fac1_mc = -grav / (cpd * theta_ref_mc**2) * inv_ddqz_z_full return d2dexdz2_fac1_mc @@ -379,7 +387,7 @@ def _compute_d2dexdz2_fac2_mc( h_scal_bg: wpfloat, igradp_method: int32, ) -> Field[[CellDim, KDim], vpfloat]: - if igradp_method <= int32(3): + if igradp_method <= 3: d2dexdz2_fac2_mc = ( 2.0 * grav @@ -675,6 +683,8 @@ def compute_vwind_impl_wgt( z_ddxn_z_half_e: intermediate storage for field z_ddxt_z_half_e: intermediate storage for field dual_edge_length: dual_edge_length + vct_a: Field[[KDim], float] + z_ifc: geometric height on half levels vwind_impl_wgt: (output) offcentering in vertical mass flux vwind_offctr: off-centering in vertical wind solver horizontal_start: horizontal start index diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfac_c.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfac_c.py index ed8259f6c1..b7375a4b43 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfac_c.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfac_c.py @@ -50,8 +50,8 @@ def _compute_wgtfac_c( k: Field[[KDim], int32], nlev: int32, ) -> Field[[CellDim, KDim], wpfloat]: - wgt_fac_c = where((k > int32(0)) & (k < nlev), _compute_wgtfac_c_inner(z_ifc), z_ifc) - wgt_fac_c = where(k == int32(0), _compute_wgtfac_c_0(z_ifc=z_ifc), wgt_fac_c) + wgt_fac_c = where((k > 0) & (k < nlev), _compute_wgtfac_c_inner(z_ifc), z_ifc) + wgt_fac_c = where(k == 0, _compute_wgtfac_c_0(z_ifc=z_ifc), wgt_fac_c) wgt_fac_c = where(k == nlev, _compute_wgtfac_c_nlev(z_ifc=z_ifc), wgt_fac_c) return wgt_fac_c diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq.py new file mode 100644 index 0000000000..cf0be286ed --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq.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 + + +def compute_z1_z2_z3(z_ifc, i1, i2, i3, i4): + z1 = 0.5 * (z_ifc[:, i2] - z_ifc[:, i1]) + z2 = 0.5 * (z_ifc[:, i2] + z_ifc[:, i3]) - z_ifc[:, i1] + z3 = 0.5 * (z_ifc[:, i3] + z_ifc[:, i4]) - z_ifc[:, i1] + return z1, z2, z3 + + +def compute_wgtfacq_c_dsl( + z_ifc: np.array, + nlev: int, +) -> np.array: + """ + Compute weighting factor for quadratic interpolation to surface. + + Args: + z_ifc: Field[CellDim, KDim] (half levels), geometric height at the vertical interface of cells. + nlev: int, last k level + Returns: + Field[CellDim, KDim] (full levels) + """ + wgtfacq_c = np.zeros((z_ifc.shape[0], nlev + 1)) + wgtfacq_c_dsl = np.zeros((z_ifc.shape[0], nlev)) + z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) + + wgtfacq_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) + wgtfacq_c[:, 1] = (z1 - wgtfacq_c[:, 2] * (z1 - z3)) / (z1 - z2) + wgtfacq_c[:, 0] = 1.0 - (wgtfacq_c[:, 1] + wgtfacq_c[:, 2]) + + wgtfacq_c_dsl[:, nlev - 1] = wgtfacq_c[:, 0] + wgtfacq_c_dsl[:, nlev - 2] = wgtfacq_c[:, 1] + wgtfacq_c_dsl[:, nlev - 3] = wgtfacq_c[:, 2] + + return wgtfacq_c_dsl + + +def compute_wgtfacq_e_dsl( + e2c, + z_ifc: np.array, + z_aux_c: np.array, + c_lin_e: np.array, + n_edges: int, + nlev: int, +): + """ + Compute weighting factor for quadratic interpolation to surface. + + Args: + e2c: Edge to Cell offset + z_ifc: geometric height at the vertical interface of cells. + z_aux_c: interpolation of weighting coefficients to edges + c_lin_e: interpolation field + n_edges: number of edges + nlev: int, last k level + Returns: + Field[EdgeDim, KDim] (full levels) + """ + wgtfacq_e_dsl = np.zeros(shape=(n_edges, nlev + 1)) + z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) + wgtfacq_c_dsl = compute_wgtfacq_c_dsl(z_ifc, nlev) + z_aux_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) + z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, nlev - 3] * (z1 - z3)) / (z1 - z2) + z_aux_c[:, 0] = 1.0 - (wgtfacq_c_dsl[:, nlev - 2] + wgtfacq_c_dsl[:, nlev - 3]) + + z1, z2, z3 = compute_z1_z2_z3(z_ifc, 0, 1, 2, 3) + z_aux_c[:, 5] = z1 * z2 / (z2 - z3) / (z1 - z3) + z_aux_c[:, 4] = (z1 - z_aux_c[:, 5] * (z1 - z3)) / (z1 - z2) + z_aux_c[:, 3] = 1.0 - (z_aux_c[:, 4] + z_aux_c[:, 5]) + + c_lin_e = c_lin_e[:, :, np.newaxis] + z_aux_e = np.sum(c_lin_e * z_aux_c[e2c], axis=1) + + wgtfacq_e_dsl[:, nlev] = z_aux_e[:, 0] + wgtfacq_e_dsl[:, nlev - 1] = z_aux_e[:, 1] + wgtfacq_e_dsl[:, nlev - 2] = z_aux_e[:, 2] + + return wgtfacq_e_dsl diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py index f78f50f188..a6d2d236c9 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py @@ -11,39 +11,3 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np - - -def compute_z1_z2_z3(z_ifc, i1, i2, i3, i4): - z1 = 0.5 * (z_ifc[:, i2] - z_ifc[:, i1]) - z2 = 0.5 * (z_ifc[:, i2] + z_ifc[:, i3]) - z_ifc[:, i1] - z3 = 0.5 * (z_ifc[:, i3] + z_ifc[:, i4]) - z_ifc[:, i1] - return z1, z2, z3 - - -def compute_wgtfacq_c_dsl( - z_ifc: np.array, - nlev: int, -) -> np.array: - """ - Compute weighting factor for quadratic interpolation to surface. - - Args: - z_ifc: Field[CellDim, KDim] (half levels), geometric height at the vertical interface of cells. - nlev: int, last k level - Returns: - Field[CellDim, KDim] (full levels) - """ - wgtfacq_c = np.zeros((z_ifc.shape[0], nlev + 1)) - wgtfacq_c_dsl = np.zeros((z_ifc.shape[0], nlev)) - z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) - - wgtfacq_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) - wgtfacq_c[:, 1] = (z1 - wgtfacq_c[:, 2] * (z1 - z3)) / (z1 - z2) - wgtfacq_c[:, 0] = 1.0 - (wgtfacq_c[:, 1] + wgtfacq_c[:, 2]) - - wgtfacq_c_dsl[:, nlev - 1] = wgtfacq_c[:, 0] - wgtfacq_c_dsl[:, nlev - 2] = wgtfacq_c[:, 1] - wgtfacq_c_dsl[:, nlev - 3] = wgtfacq_c[:, 2] - - return wgtfacq_c_dsl diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 470d30503a..1b6f757b41 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -239,7 +239,7 @@ class StencilTest: """ PROGRAM: ClassVar[Program] - OUTPUTS: ClassVar[tuple[str, ...]] + OUTPUTS: ClassVar[tuple[str | Output, ...]] def __init_subclass__(cls, **kwargs): # Add two methods for verification and benchmarking. In order to have names that 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 index 21f292bf02..6b38dc3a88 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -539,12 +539,14 @@ def rbf_vec_coeff_e(self): ).transpose() return as_field((EdgeDim, E2C2EDim), buffer) + @IconSavepoint.optionally_registered() def rbf_vec_coeff_c1(self): buffer = np.squeeze( self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float) ).transpose() return as_field((CellDim, C2E2C2EDim), buffer) + @IconSavepoint.optionally_registered() def rbf_vec_coeff_c2(self): buffer = np.squeeze( self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float) @@ -583,6 +585,7 @@ def hmask_dd3d(self): def inv_ddqz_z_full(self): return self._get_field("inv_ddqz_z_full", CellDim, KDim) + @IconSavepoint.optionally_registered(CellDim, KDim) def ddqz_z_full(self): return self._get_field("ddqz_z_full", CellDim, KDim) diff --git a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py index aaa8981362..6227f2d00c 100644 --- a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py +++ b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py @@ -25,7 +25,6 @@ import numpy as np import pytest -from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import EdgeDim from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex, RefinCtrlLevel @@ -56,7 +55,7 @@ def test_compute_nudgecoeffs_e( grf_nudge_start_e = RefinCtrlLevel.boundary_nudging_start(EdgeDim) nudge_max_coeff = wpfloat(0.375) nudge_efold_width = wpfloat(2.0) - nudge_zone_width = int32(10) + nudge_zone_width = 10 horizontal_start = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.nudging_2nd_level(EdgeDim) diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq.py b/model/common/tests/metric_tests/test_compute_wgtfacq.py index 8d5194cd61..c38afa6262 100644 --- a/model/common/tests/metric_tests/test_compute_wgtfacq.py +++ b/model/common/tests/metric_tests/test_compute_wgtfacq.py @@ -25,10 +25,12 @@ import pytest -from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim -from icon4py.model.common.metrics.stencils.compute_wgtfacq_c_dsl import compute_wgtfacq_c_dsl -from icon4py.model.common.metrics.stencils.compute_wgtfacq_e_dsl import compute_wgtfacq_e_dsl -from icon4py.model.common.test_utils.helpers import dallclose, zero_field +from icon4py.model.common.dimension import E2CDim +from icon4py.model.common.metrics.stencils.compute_wgtfacq import ( + compute_wgtfacq_c_dsl, + compute_wgtfacq_e_dsl, +) +from icon4py.model.common.test_utils.helpers import dallclose @pytest.mark.datatest @@ -44,7 +46,6 @@ def test_compute_wgtfacq_c_dsl(icon_grid, metrics_savepoint): @pytest.mark.datatest def test_compute_wgtfacq_e_dsl(metrics_savepoint, interpolation_savepoint, icon_grid): - wgtfacq_e_dsl_full = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) wgtfacq_e_dsl_ref = metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels + 1) wgtfacq_e_dsl_full = compute_wgtfacq_e_dsl( @@ -52,7 +53,7 @@ def test_compute_wgtfacq_e_dsl(metrics_savepoint, interpolation_savepoint, icon_ z_ifc=metrics_savepoint.z_ifc().asnumpy(), z_aux_c=metrics_savepoint.wgtfac_c().asnumpy(), c_lin_e=interpolation_savepoint.c_lin_e().asnumpy(), - wgtfacq_e_dsl=wgtfacq_e_dsl_full.asnumpy(), + n_edges=icon_grid.num_edges, nlev=icon_grid.num_levels, ) diff --git a/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py b/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py new file mode 100644 index 0000000000..397a8e0c3a --- /dev/null +++ b/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py @@ -0,0 +1,116 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file 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 gt4py.next import as_field +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.common.dimension import ( + CellDim, + E2CDim, + ECDim, + EdgeDim, + KDim, +) +from icon4py.model.common.grid.horizontal import ( + HorizontalMarkerIndex, +) +from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( + _cell_2_edge_interpolation, +) +from icon4py.model.common.metrics.metric_fields import ( + _compute_flat_idx, + _compute_z_aux2, + compute_z_mc, +) +from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl +from icon4py.model.common.test_utils.helpers import ( + dallclose, + flatten_first_two_dims, + is_roundtrip, + zero_field, +) + + +@pytest.mark.datatest +def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") + zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() + z_mc = zero_field(icon_grid, CellDim, KDim) + z_ifc = metrics_savepoint.z_ifc() + k_lev = as_field((KDim,), np.arange(icon_grid.num_levels, dtype=int)) + z_me = zero_field(icon_grid, EdgeDim, KDim) + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + end_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) + compute_z_mc.with_backend(backend)( + z_ifc, + z_mc, + horizontal_start=int32(0), + horizontal_end=icon_grid.num_cells, + vertical_start=int32(0), + vertical_end=int32(icon_grid.num_levels), + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + _cell_2_edge_interpolation( + in_field=z_mc, + coeff=interpolation_savepoint.c_lin_e(), + out=z_me, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + flat_idx = zero_field(icon_grid, EdgeDim, KDim) + _compute_flat_idx( + z_me=z_me, + z_ifc=z_ifc, + k_lev=k_lev, + out=flat_idx, + domain={ + EdgeDim: (horizontal_start_edge, icon_grid.num_edges), + KDim: (int32(0), icon_grid.num_levels), + }, + offset_provider={ + "E2C": icon_grid.get_offset_provider("E2C"), + "Koff": icon_grid.get_offset_provider("Koff"), + }, + ) + flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) + z_ifc_sliced = as_field((CellDim,), z_ifc.asnumpy()[:, icon_grid.num_levels]) + z_aux2 = zero_field(icon_grid, EdgeDim) + _compute_z_aux2( + z_ifc=z_ifc_sliced, + out=z_aux2, + domain={EdgeDim: (end_edge_nudging, icon_grid.num_edges)}, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + + zdiff_gradp_full_np = compute_zdiff_gradp_dsl( + e2c=icon_grid.connectivities[E2CDim], + z_me=z_me.asnumpy(), + z_mc=z_mc.asnumpy(), + z_ifc=metrics_savepoint.z_ifc().asnumpy(), + flat_idx=flat_idx_np, + z_aux2=z_aux2.asnumpy(), + nlev=icon_grid.num_levels, + horizontal_start=horizontal_start_edge, + horizontal_start_1=end_edge_nudging, + nedges=icon_grid.num_edges, + ) + zdiff_gradp_full_field = flatten_first_two_dims( + ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np) + ) + assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 70bd4e4d99..b121d60e4c 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -23,14 +23,14 @@ C2E2CDim, CECDim, CellDim, - E2CDim, - ECDim, EdgeDim, KDim, V2CDim, VertexDim, ) from icon4py.model.common.grid.horizontal import ( + _GRF_NUDGEZONE_START_EDGES, + _GRF_NUDGEZONE_WIDTH, HorizontalMarkerIndex, _compute_cells2verts, ) @@ -40,6 +40,7 @@ ) from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( + MetricsConfig, _compute_flat_idx, _compute_max_nbhgt, _compute_maxslp_maxhgtd, @@ -71,7 +72,6 @@ _compute_k_start_end, compute_diffusion_metrics, ) -from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, @@ -106,9 +106,9 @@ def reference( def input_data(self, grid) -> dict: z_mc = zero_field(grid, CellDim, KDim) z_if = random_field(grid, CellDim, KDim, extend={KDim: 1}) - horizontal_start = int32(0) + horizontal_start = 0 horizontal_end = grid.num_cells - vertical_start = int32(0) + vertical_start = 0 vertical_end = grid.num_levels return dict( @@ -133,9 +133,9 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): compute_z_mc.with_backend(backend)( z_ifc, z_mc, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -170,9 +170,9 @@ def test_compute_ddqz_z_full_and_inverse(icon_grid, metrics_savepoint, backend): z_ifc=z_ifc, ddqz_z_full=ddqz_z_full, inv_ddqz_z_full=inv_ddqz_z_full, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), + vertical_start=0, vertical_end=icon_grid.num_levels, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -195,7 +195,7 @@ def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, back divdamp_trans_start=divdamp_trans_start, divdamp_trans_end=divdamp_trans_end, divdamp_type=divdamp_type, - vertical_start=int32(0), + vertical_start=0, vertical_end=icon_grid.num_levels, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -222,7 +222,7 @@ def test_compute_rayleigh_w(icon_grid, metrics_savepoint, grid_savepoint, backen rayleigh_coeff=rayleigh_coeff, vct_a_1=vct_a_1, pi_const=math.pi, - vertical_start=int32(0), + vertical_start=0, vertical_end=grid_savepoint.nrdmax().item() + 1, offset_provider={}, ) @@ -246,9 +246,9 @@ def test_compute_coeff_dwdz(icon_grid, metrics_savepoint, grid_savepoint, backen z_ifc=metrics_savepoint.z_ifc(), coeff1_dwdz=coeff1_dwdz_full, coeff2_dwdz=coeff2_dwdz_full, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(1), + vertical_start=1, vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -266,9 +266,9 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b compute_z_mc.with_backend(backend)( z_ifc=z_ifc, z_mc=z_mc, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -294,10 +294,10 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b grav=grav, del_t_bg=del_t_bg, h_scal_bg=h_scal_bg, - igradp_method=int32(3), - horizontal_start=int32(0), + igradp_method=3, + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -375,7 +375,7 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): compute_vwind_expl_wgt.with_backend(backend)( vwind_impl_wgt=vwind_impl_wgt, vwind_expl_wgt=vwind_expl_wgt_full, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) @@ -396,8 +396,6 @@ def test_compute_ddqz_z_full_e( vertical_start = 0 vertical_end = icon_grid.num_levels ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) - # TODO: perhaps write a program with ddqz_z_full_e name and call fieldop _cells2edges... from there - # TODO: This way it's clear where this field is computed and we cna more easily avoid duplicates cell_2_edge_interpolation.with_backend(backend)( in_field=ddqz_z_full, coeff=c_lin_e, @@ -522,13 +520,13 @@ def test_compute_exner_exfac( horizontal_start = icon_grid.get_start_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1 ) - exner_exfac = constant_field(icon_grid, constants.exner_expol, CellDim, KDim) + exner_exfac = constant_field(icon_grid, MetricsConfig.exner_expol, CellDim, KDim) exner_exfac_ref = metrics_savepoint.exner_exfac() compute_exner_exfac.with_backend(backend)( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), exner_exfac=exner_exfac, - exner_expol=constants.exner_expol, + exner_expol=MetricsConfig.exner_expol, horizontal_start=horizontal_start, horizontal_end=icon_grid.num_cells, vertical_start=int32(0), @@ -539,79 +537,6 @@ def test_compute_exner_exfac( assert dallclose(exner_exfac.asnumpy(), exner_exfac_ref.asnumpy(), rtol=1.0e-10) -# TODO -@pytest.mark.datatest -def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() - z_mc = zero_field(icon_grid, CellDim, KDim) - z_ifc = metrics_savepoint.z_ifc() - k_lev = as_field((KDim,), np.arange(icon_grid.num_levels, dtype=int)) - z_me = zero_field(icon_grid, EdgeDim, KDim) - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - end_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) - compute_z_mc.with_backend(backend)( - z_ifc, - z_mc, - horizontal_start=int32(0), - horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), - vertical_end=int32(icon_grid.num_levels), - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - _cell_2_edge_interpolation( - in_field=z_mc, - coeff=interpolation_savepoint.c_lin_e(), - out=z_me, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - flat_idx = zero_field(icon_grid, EdgeDim, KDim) - _compute_flat_idx( - z_me=z_me, - z_ifc=z_ifc, - k_lev=k_lev, - out=flat_idx, - domain={ - EdgeDim: (horizontal_start_edge, icon_grid.num_edges), - KDim: (int32(0), icon_grid.num_levels), - }, - offset_provider={ - "E2C": icon_grid.get_offset_provider("E2C"), - "Koff": icon_grid.get_offset_provider("Koff"), - }, - ) - flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) - z_ifc_sliced = as_field((CellDim,), z_ifc.asnumpy()[:, icon_grid.num_levels]) - z_aux2 = zero_field(icon_grid, EdgeDim) - _compute_z_aux2( - z_ifc=z_ifc_sliced, - out=z_aux2, - domain={EdgeDim: (end_edge_nudging, icon_grid.num_edges)}, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - - zdiff_gradp_full_np = compute_zdiff_gradp_dsl( - e2c=icon_grid.connectivities[E2CDim], - z_me=z_me.asnumpy(), - z_mc=z_mc.asnumpy(), - z_ifc=metrics_savepoint.z_ifc().asnumpy(), - flat_idx=flat_idx_np, - z_aux2=z_aux2.asnumpy(), - nlev=icon_grid.num_levels, - horizontal_start=horizontal_start_edge, - horizontal_start_1=end_edge_nudging, - nedges=icon_grid.num_edges, - ) - zdiff_gradp_full_field = flatten_first_two_dims( - ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np) - ) - assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) - - @pytest.mark.datatest def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend @@ -892,8 +817,8 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen compute_hmask_dd3d( e_refin_ctrl=e_refin_ctrl, hmask_dd3d=hmask_dd3d_full, - grf_nudge_start_e=int32(constants.grf_nudge_start_e), - grf_nudgezone_width=int32(constants.grf_nudgezone_width), + grf_nudge_start_e=int32(_GRF_NUDGEZONE_START_EDGES), + grf_nudgezone_width=int32(_GRF_NUDGEZONE_WIDTH), horizontal_start=horizontal_start, horizontal_end=icon_grid.num_edges, offset_provider={}, @@ -906,7 +831,8 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() @@ -1029,4 +955,4 @@ def test_compute_diffusion_metrics( assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy()) - assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef()) + assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef().asnumpy()) diff --git a/model/common/tests/metric_tests/test_reference_atmosphere.py b/model/common/tests/metric_tests/test_reference_atmosphere.py index 78fb3c048c..de0224434d 100644 --- a/model/common/tests/metric_tests/test_reference_atmosphere.py +++ b/model/common/tests/metric_tests/test_reference_atmosphere.py @@ -49,7 +49,7 @@ def test_compute_reference_atmosphere_fields_on_full_level_masspoints( rho_ref_mc = zero_field(icon_grid, CellDim, KDim, dtype=wpfloat) theta_ref_mc = zero_field(icon_grid, CellDim, KDim, dtype=wpfloat) z_mc = zero_field(icon_grid, CellDim, KDim, dtype=wpfloat) - start = int32(0) + start = 0 horizontal_end = icon_grid.num_cells vertical_end = icon_grid.num_levels compute_z_mc.with_backend(backend)( @@ -99,7 +99,7 @@ def test_compute_reference_atmsophere_on_half_level_mass_points( exner_ref_ic = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}, dtype=wpfloat) rho_ref_ic = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}, dtype=wpfloat) theta_ref_ic = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}, dtype=wpfloat) - start = int32(0) + start = 0 horizontal_end = icon_grid.num_cells vertical_end = icon_grid.num_levels + 1 @@ -164,7 +164,7 @@ def test_compute_reference_atmosphere_on_full_level_edge_fields( ) num_cells = int32(icon_grid.num_cells) num_edges = int(icon_grid.num_edges) - vertical_start = int32(0) + vertical_start = 0 vertical_end = int32(icon_grid.num_levels) compute_z_mc.with_backend(backend)( z_ifc=z_ifc, diff --git a/model/common/tests/stencil_tests/test_cell_2_edge_interpolation.py b/model/common/tests/stencil_tests/test_cell_2_edge_interpolation.py index c1242a0fd2..2fc08cd0e3 100644 --- a/model/common/tests/stencil_tests/test_cell_2_edge_interpolation.py +++ b/model/common/tests/stencil_tests/test_cell_2_edge_interpolation.py @@ -47,8 +47,8 @@ def input_data(self, grid): in_field=in_field, coeff=coeff, out_field=out_field, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/common/tests/stencil_tests/test_diagnose_pressure.py b/model/common/tests/stencil_tests/test_diagnose_pressure.py index 20204207af..9a21c9eeb5 100644 --- a/model/common/tests/stencil_tests/test_diagnose_pressure.py +++ b/model/common/tests/stencil_tests/test_diagnose_pressure.py @@ -77,8 +77,8 @@ def input_data(self, grid): pressure=pressure, pressure_ifc=pressure_ifc, grav_o_rd=GRAV_O_RD, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py b/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py index 5e075f6232..e95e7e597a 100644 --- a/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py +++ b/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py @@ -65,7 +65,7 @@ def input_data(self, grid): cpd_o_rd=CPD_O_RD, p0ref=P0REF, grav_o_rd=GRAV_O_RD, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), vertical_start=int32(grid.num_levels), vertical_end=int32(grid.num_levels + 1), diff --git a/model/common/tests/stencil_tests/test_diagnose_temperature.py b/model/common/tests/stencil_tests/test_diagnose_temperature.py index f91729915e..9ebac04c24 100644 --- a/model/common/tests/stencil_tests/test_diagnose_temperature.py +++ b/model/common/tests/stencil_tests/test_diagnose_temperature.py @@ -44,8 +44,8 @@ def input_data(self, grid): theta_v=theta_v, exner=exner, temperature=temperature, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/common/tests/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py b/model/common/tests/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py index 8967cc09a3..d9f77c7abb 100644 --- a/model/common/tests/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py +++ b/model/common/tests/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py @@ -53,8 +53,8 @@ def input_data(self, grid): ptr_coeff_2=ptr_coeff_2, p_v_out=p_v_out, p_u_out=p_u_out, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 41de2ea90a..ea5c6fd17a 100644 --- a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -54,8 +54,8 @@ def input_data(self, grid): ptr_coeff_2=ptr_coeff_2, p_v_out=p_v_out, p_u_out=p_u_out, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=int32(grid.num_vertices), - vertical_start=int32(0), + vertical_start=0, vertical_end=int32(grid.num_levels), ) diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml index ae3d7af30e..f09391e41b 100644 --- a/model/driver/pyproject.toml +++ b/model/driver/pyproject.toml @@ -25,7 +25,8 @@ dependencies = [ "gt4py>=1.0.1", "icon4py-common>=0.0.5", "icon4py-atmosphere-dycore>=0.0.5", - "icon4py-atmosphere-diffusion>=0.0.5" + "icon4py-atmosphere-diffusion>=0.0.5", + "nanobind<2.0.0" ] description = "ICON model driver." dynamic = ['version'] diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index 25d75230e7..6c183c9817 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -186,10 +186,10 @@ def time_integration( for time_step in range(self._n_time_steps): log.info(f"simulation date : {self._simulation_date} run timestep : {time_step}") log.info( - f" MAX VN: {prognostic_state_list[self._now].vn.asnumpy().max():.5e} , MAX W: {prognostic_state_list[self._now].w.asnumpy().max():.5e}" + f" MAX VN: {prognostic_state_list[self._now].vn.asnumpy().max():.15e} , MAX W: {prognostic_state_list[self._now].w.asnumpy().max():.15e}" ) log.info( - f" MAX RHO: {prognostic_state_list[self._now].rho.asnumpy().max():.5e} , MAX THETA_V: {prognostic_state_list[self._now].theta_v.asnumpy().max():.5e}" + f" MAX RHO: {prognostic_state_list[self._now].rho.asnumpy().max():.15e} , MAX THETA_V: {prognostic_state_list[self._now].theta_v.asnumpy().max():.15e}" ) # TODO (Chia Rui): check with Anurag about printing of max and min of variables. @@ -294,6 +294,8 @@ def initialize( props: ProcessProperties, serialization_type: SerializationType, experiment_type: ExperimentType, + grid_root, + grid_level, ): """ Inititalize the driver run. @@ -323,16 +325,24 @@ def initialize( log.info(f"reading configuration: experiment {experiment_type}") config = read_config(experiment_type) - decomp_info = read_decomp_info(file_path, props, serialization_type) + decomp_info = read_decomp_info(file_path, props, serialization_type, grid_root, grid_level) log.info(f"initializing the grid from '{file_path}'") - icon_grid = read_icon_grid(file_path, rank=props.rank, ser_type=serialization_type) + icon_grid = read_icon_grid( + file_path, + rank=props.rank, + ser_type=serialization_type, + grid_root=grid_root, + grid_level=grid_level, + ) log.info(f"reading input fields from '{file_path}'") (edge_geometry, cell_geometry, vertical_geometry, c_owner_mask) = read_geometry_fields( file_path, damping_height=config.run_config.damping_height, rank=props.rank, ser_type=serialization_type, + grid_root=grid_root, + grid_level=grid_level, ) ( diffusion_metric_state, @@ -341,7 +351,11 @@ def initialize( solve_nonhydro_interpolation_state, diagnostic_metric_state, ) = read_static_fields( - file_path, rank=props.rank, ser_type=serialization_type, experiment_type=experiment_type + file_path, + rank=props.rank, + ser_type=serialization_type, + grid_root=grid_root, + grid_level=grid_level, ) log.info("initializing diffusion") @@ -418,7 +432,9 @@ def initialize( help="serialization type for grid info and static fields", ) @click.option("--experiment_type", default="any", help="experiment selection") -def main(input_path, run_path, mpi, serialization_type, experiment_type): +@click.option("--grid_root", default=2, help="experiment selection") +@click.option("--grid_level", default=4, help="experiment selection") +def main(input_path, run_path, mpi, serialization_type, experiment_type, grid_root, grid_level): """ Run the driver. @@ -449,7 +465,9 @@ def main(input_path, run_path, mpi, serialization_type, experiment_type): diagnostic_state, prep_adv, inital_divdamp_fac_o2, - ) = initialize(Path(input_path), parallel_props, serialization_type, experiment_type) + ) = initialize( + Path(input_path), parallel_props, serialization_type, experiment_type, grid_root, grid_level + ) log.info(f"Starting ICON dycore run: {timeloop.simulation_date.isoformat()}") log.info( f"input args: input_path={input_path}, n_time_steps={timeloop.n_time_steps}, ending date={timeloop.run_config.end_date}" diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index c6a8711c08..968faf973e 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -25,6 +25,7 @@ DiffusionInterpolationState, DiffusionMetricState, ) +from icon4py.model.atmosphere.dycore.init_exner_pr import init_exner_pr from icon4py.model.atmosphere.dycore.state_utils.states import ( DiagnosticStateNonHydro, InterpolationState, @@ -45,7 +46,6 @@ from icon4py.model.common.decomposition.definitions import DecompositionInfo, ProcessProperties from icon4py.model.common.decomposition.mpi_decomposition import ParallelLogger from icon4py.model.common.dimension import ( - C2E2C2EDim, CEDim, CellDim, EdgeDim, @@ -88,29 +88,35 @@ class SerializationType(str, Enum): class ExperimentType(str, Enum): - """Jablonowski-Williamson test""" - JABW = "jabw" - """any test with initial conditions read from serialized data""" + """initial condition of Jablonowski-Williamson test""" ANY = "any" + """any test with initial conditions read from serialized data (remember to set correct SIMULATION_START_DATE)""" def read_icon_grid( - path: Path, rank=0, ser_type: SerializationType = SerializationType.SB + path: Path, + rank=0, + ser_type: SerializationType = SerializationType.SB, + grid_root=2, + grid_level=4, ) -> IconGrid: """ Read icon grid. Args: path: path where to find the input data + rank: mpi rank of the current compute node ser_type: type of input data. Currently only 'sb (serialbox)' is supported. It reads from ppser serialized test data + grid_root: global grid root division number + grid_level: global grid refinement number 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(2, 4) + .from_savepoint_grid(grid_root, grid_level) .construct_icon_grid(on_gpu=False) ) else: @@ -123,7 +129,29 @@ def model_initialization_jabw( edge_param: EdgeParams, path: Path, rank=0, -): +) -> tuple[ + DiffusionDiagnosticState, + DiagnosticStateNonHydro, + PrepAdvection, + float, + DiagnosticState, + PrognosticState, + PrognosticState, +]: + """ + Initial condition of Jablonowski-Williamson test. Set jw_up to values larger than 0.01 if + you want to run baroclinic case. + + Args: + icon_grid: IconGrid + cell_param: cell properties + edge_param: edge properties + path: path where to find the input data + rank: mpi rank of the current compute node + Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, + PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic + variables (now and next). + """ data_provider = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) @@ -152,6 +180,9 @@ def model_initialization_jabw( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1 ) grid_idx_edge_end = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.end(EdgeDim)) + grid_idx_cell_interior_start = icon_grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) + ) grid_idx_cell_start_plus1 = icon_grid.get_end_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1 ) @@ -324,6 +355,19 @@ def model_initialization_jabw( log.info("U, V computation completed.") + exner_pr = _allocate(CellDim, KDim, grid=icon_grid) + init_exner_pr( + exner, + data_provider.from_metrics_savepoint().exner_ref_mc(), + exner_pr, + grid_idx_cell_interior_start, + grid_idx_cell_end, + 0, + icon_grid.num_levels, + offset_provider={}, + ) + log.info("exner_pr initialization completed.") + diagnostic_state = DiagnosticState( pressure=pressure, pressure_ifc=pressure_ifc, @@ -355,7 +399,7 @@ def model_initialization_jabw( ) solve_nonhydro_diagnostic_state = DiagnosticStateNonHydro( theta_v_ic=_allocate(CellDim, KDim, grid=icon_grid, is_halfdim=True), - exner_pr=_allocate(CellDim, KDim, grid=icon_grid), + exner_pr=exner_pr, rho_ic=_allocate(CellDim, KDim, grid=icon_grid, is_halfdim=True), ddt_exner_phy=_allocate(CellDim, KDim, grid=icon_grid), grf_tend_rho=_allocate(CellDim, KDim, grid=icon_grid), @@ -396,19 +440,28 @@ def model_initialization_jabw( ) -def model_initialization_serialbox(icon_grid: IconGrid, path: Path, rank=0): +def model_initialization_serialbox( + icon_grid: IconGrid, path: Path, rank=0 +) -> tuple[ + DiffusionDiagnosticState, + DiagnosticStateNonHydro, + PrepAdvection, + float, + DiagnosticState, + PrognosticState, + PrognosticState, +]: """ - Read prognostic and diagnostic state from serialized data. + Initial condition read from serialized data. Diagnostic variables are allocated as zero + fields. Args: - icon_grid: icon grid - path: path to the serialized input data + icon_grid: IconGrid + path: path where to find the input data rank: mpi rank of the current compute node - - 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 - + Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, + PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic + variables (now and next). """ data_provider = sb.IconSerialDataProvider( @@ -453,9 +506,8 @@ def model_initialization_serialbox(icon_grid: IconGrid, path: Path, rank=0): diagnostic_state = DiagnosticState( pressure=_allocate(CellDim, KDim, grid=icon_grid), - pressure_ifc=_allocate(CellDim, KDim, grid=icon_grid), + pressure_ifc=_allocate(CellDim, KDim, grid=icon_grid, is_halfdim=True), temperature=_allocate(CellDim, KDim, grid=icon_grid), - pressure_sfc=_allocate(CellDim, grid=icon_grid), u=_allocate(CellDim, KDim, grid=icon_grid), v=_allocate(CellDim, KDim, grid=icon_grid), ) @@ -501,6 +553,21 @@ def read_initial_state( PrognosticState, PrognosticState, ]: + """ + Read initial prognostic and diagnostic fields. + + Args: + icon_grid: IconGrid + cell_param: cell properties + edge_param: edge properties + path: path to the serialized input data + rank: mpi rank of the current compute node + experiment_type: (optional) defaults to ANY=any, type of initial condition to be read + + Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, + PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic + variables (now and next). + """ if experiment_type == ExperimentType.JABW: ( diffusion_diagnostic_state, @@ -536,16 +603,23 @@ def read_initial_state( def read_geometry_fields( - path: Path, damping_height, rank=0, ser_type: SerializationType = SerializationType.SB + path: Path, + damping_height, + rank=0, + ser_type: SerializationType = SerializationType.SB, + grid_root=2, + grid_level=4, ) -> tuple[EdgeParams, CellParams, VerticalModelParams, Field[[CellDim], bool]]: """ Read fields containing grid properties. Args: path: path to the serialized input data - damping_height: damping height for Rayleigh and divergence damping TODO (CHia Rui): Check - rank: + damping_height: damping height for Rayleigh and divergence damping + rank: mpi rank of the current compute node ser_type: (optional) defaults to SB=serialbox, type of input data to be read + grid_root: global grid root division number + grid_level: global grid refinement number 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. @@ -553,7 +627,7 @@ def read_geometry_fields( if ser_type == SerializationType.SB: sp = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank - ).from_savepoint_grid(2, 4) + ).from_savepoint_grid(grid_root, grid_level) edge_geometry = sp.construct_edge_geometry() cell_geometry = sp.construct_cell_geometry() vertical_geometry = VerticalModelParams( @@ -571,12 +645,14 @@ def read_decomp_info( path: Path, procs_props: ProcessProperties, ser_type=SerializationType.SB, + grid_root=2, + grid_level=4, ) -> DecompositionInfo: if ser_type == SerializationType.SB: sp = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), True, procs_props.rank ) - return sp.from_savepoint_grid(2, 4).construct_decomposition_info() + return sp.from_savepoint_grid(grid_root, grid_level).construct_decomposition_info() else: raise NotImplementedError(SB_ONLY_MSG) @@ -585,7 +661,8 @@ def read_static_fields( path: Path, rank=0, ser_type: SerializationType = SerializationType.SB, - experiment_type: ExperimentType = ExperimentType.JABW, + grid_root=2, + grid_level=4, ) -> tuple[ DiffusionMetricState, DiffusionInterpolationState, @@ -600,7 +677,8 @@ def read_static_fields( 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 - experiment_type: TODO (CHia RUi): Add description + grid_root: global grid root division number + grid_level: global grid refinement number Returns: a tuple containing the metric_state and interpolation state, @@ -613,7 +691,7 @@ def read_static_fields( ) icon_grid = ( sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) - .from_savepoint_grid(2, 4) + .from_savepoint_grid(grid_root, grid_level) .construct_icon_grid(on_gpu=False) ) diffusion_interpolation_state = construct_interpolation_state_for_diffusion( @@ -679,19 +757,11 @@ def read_static_fields( coeff_gradekin=metrics_savepoint.coeff_gradekin(), ) - if experiment_type == ExperimentType.JABW: - diagnostic_metric_state = DiagnosticMetricState( - ddqz_z_full=metrics_savepoint.ddqz_z_full(), - rbf_vec_coeff_c1=data_provider.from_interpolation_savepoint().rbf_vec_coeff_c1(), - rbf_vec_coeff_c2=data_provider.from_interpolation_savepoint().rbf_vec_coeff_c2(), - ) - else: - # ddqz_z_full is not serialized for mch_ch_r04b09_dsl and exclaim_ape_R02B04 - diagnostic_metric_state = DiagnosticMetricState( - ddqz_z_full=_allocate(CellDim, KDim, grid=icon_grid, dtype=float), - rbf_vec_coeff_c1=_allocate(CellDim, C2E2C2EDim, grid=icon_grid, dtype=float), - rbf_vec_coeff_c2=_allocate(CellDim, C2E2C2EDim, grid=icon_grid, dtype=float), - ) + diagnostic_metric_state = DiagnosticMetricState( + ddqz_z_full=metrics_savepoint.ddqz_z_full(), + rbf_vec_coeff_c1=interpolation_savepoint.rbf_vec_coeff_c1(), + rbf_vec_coeff_c2=interpolation_savepoint.rbf_vec_coeff_c2(), + ) return ( diffusion_metric_state, diff --git a/model/driver/src/icon4py/model/driver/jablonowski_willamson_testcase.py b/model/driver/src/icon4py/model/driver/jablonowski_willamson_testcase.py index ddec9fee7a..599f076d88 100644 --- a/model/driver/src/icon4py/model/driver/jablonowski_willamson_testcase.py +++ b/model/driver/src/icon4py/model/driver/jablonowski_willamson_testcase.py @@ -29,7 +29,21 @@ def zonalwind_2_normalwind_jabw_numpy( primal_normal_x: np.array, eta_v_e: np.array, ): - """mask = np.repeat(np.expand_dims(mask, axis=-1), eta_v_e.shape[1], axis=1)""" + """ + Compute normal wind at edge center from vertical eta coordinate (eta_v_e). + + Args: + icon_grid: IconGrid + jw_u0: base zonal wind speed factor + jw_up: perturbation amplitude + lat_perturbation_center: perturbation center in latitude + lon_perturbation_center: perturbation center in longitude + edge_lat: edge center latitude + edge_lon: edge center longitude + primal_normal_x: zonal component of primal normal vector at edge center + eta_v_e: vertical eta coordinate at edge center + Returns: normal wind + """ mask = np.ones((icon_grid.num_edges, icon_grid.num_levels), dtype=bool) mask[ 0 : icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1), : diff --git a/model/driver/tests/test_testcase.py b/model/driver/tests/test_testcase.py index ad1211c282..341251fe17 100644 --- a/model/driver/tests/test_testcase.py +++ b/model/driver/tests/test_testcase.py @@ -89,3 +89,9 @@ def test_jabw_initial_condition( diagnostic_state.pressure_sfc.asnumpy(), data_provider.from_savepoint_jabw_init().pressure_sfc().asnumpy(), ) + + assert dallclose( + solve_nonhydro_diagnostic_state.exner_pr.asnumpy(), + data_provider.from_savepoint_jabw_diagnostic().exner_pr().asnumpy(), + atol=1.0e-14, + ) diff --git a/model/tox.ini b/model/tox.ini index 2718b464e4..39ce4ab666 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -30,17 +30,18 @@ allowlist_externals = commands = pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-skip -n auto {posargs} pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/dycore/tests/dycore_stencil_tests --benchmark-skip -n auto {posargs} + pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/advection/tests/advection_stencil_tests --benchmark-skip -n auto {posargs} [testenv:run_benchmarks] commands = pytest -s -m "not slow_tests" atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-only {posargs} pytest -s -m "not slow_tests" atmosphere/dycore/tests/dycore_stencil_tests --benchmark-only {posargs} + pytest -s -m "not slow_tests" atmosphere/advection/tests/advection_stencil_tests --benchmark-only {posargs} [testenv:run_model_tests] commands = pytest -v -s -m "not slow_tests" --datatest {posargs} - [testenv:dev] setenv = PIP_SRC = _external_src diff --git a/tools/README.md b/tools/README.md index d6288ad2f7..72132757df 100644 --- a/tools/README.md +++ b/tools/README.md @@ -148,6 +148,8 @@ In addition, other optional keyword arguments are the following: - `copies`: Takes a boolean string input, and controls whether before field copies should be made or not. If set to False only the `#ifdef __DSL_VERIFY` directive is generated. Defaults to true.

+- `optional_module`: Takes a boolean string input, and controls whether stencils is part of an optional module. Defaults to "None".

+ #### `!$DSL END STENCIL()` This directive denotes the end of a stencil. The required argument is `name`, which must match the name of the preceding `START STENCIL` directive. diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 8bb7d3df53..1221c0d319 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -26,7 +26,8 @@ dependencies = [ 'icon4py-common', 'tabulate>=0.8.9', 'fprettify>=0.3.7', - 'cffi>=1.5' + 'cffi>=1.5', + "nanobind<2.0.0" ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index a90afbc09e..a118199e4d 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pathlib +from typing import Optional import click @@ -27,6 +28,10 @@ logger = setup_logger(__name__) +def split_comma(ctx, param, value) -> Optional[tuple[str]]: + return tuple(v.strip() for v in value.split(",")) if value else None + + @click.group(invoke_without_command=True) @click.pass_context def main(ctx: click.Context) -> None: @@ -56,6 +61,11 @@ def main(ctx: click.Context) -> None: default=False, help="Adds fused or unfused stencils.", ) +@click.option( + "--optional-modules-to-enable", + callback=split_comma, + help="Specify a list of comma-separated optional DSL modules to enable.", +) @click.option( "--verification/--substitution", "-v/-s", @@ -77,10 +87,13 @@ def integrate( verification: bool, profile: bool, metadatagen: bool, + optional_modules_to_enable: Optional[tuple[str]], ) -> None: mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) - iface_gt4py = process_stencils(iface, fused) + iface_gt4py = process_stencils( + iface, fused, optional_modules_to_enable=optional_modules_to_enable + ) run_code_generation( input_path, output_path, diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 54e085d685..d669f220c6 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -53,6 +53,10 @@ TOLERANCE_ARGS = ["abs_tol", "rel_tol"] DEFAULT_DECLARE_IDENT_TYPE = "REAL(wp)" DEFAULT_DECLARE_SUFFIX = "before" +DEFAULT_STARTSTENCIL_ACC_PRESENT = "true" +DEFAULT_STARTSTENCIL_MERGECOPY = "false" +DEFAULT_STARTSTENCIL_COPIES = "true" +DEFAULT_STARTSTENCIL_OPTIONAL_MODULE = "None" logger = setup_logger(__name__) @@ -286,12 +290,13 @@ def create_stencil_data( for i, directive in enumerate(directives): named_args = parsed["content"][directive_cls.__name__][i] additional_attrs = self._pop_additional_attributes(dtype, named_args) - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + acc_present = string_to_bool( + pop_item_from_dict(named_args, "accpresent", DEFAULT_STARTSTENCIL_ACC_PRESENT) + ) stencil_name = _extract_stencil_name(named_args, directive) bounds = self._make_bounds(named_args) fields = self._make_fields(named_args, field_dimensions) fields_w_tolerance = self._update_tolerances(named_args, fields) - deserialised.append( dtype( name=stencil_name, @@ -305,14 +310,27 @@ def create_stencil_data( return deserialised def _pop_additional_attributes( - self, dtype: Type[StartStencilData | StartFusedStencilData], named_args: dict[str, Any] + self, + dtype: Type[StartStencilData | StartFusedStencilData], + named_args: dict[str, Any], ) -> dict: """Pop and return additional attributes specific to StartStencilData.""" additional_attrs = {} if dtype == StartStencilData: - mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) - copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) - additional_attrs = {"mergecopy": mergecopy, "copies": copies} + mergecopy = string_to_bool( + pop_item_from_dict(named_args, "mergecopy", DEFAULT_STARTSTENCIL_MERGECOPY) + ) + copies = string_to_bool( + pop_item_from_dict(named_args, "copies", DEFAULT_STARTSTENCIL_COPIES) + ) + optional_module = pop_item_from_dict( + named_args, "optional_module", DEFAULT_STARTSTENCIL_OPTIONAL_MODULE + ) + additional_attrs = { + "mergecopy": mergecopy, + "copies": copies, + "optional_module": optional_module, + } return additional_attrs @staticmethod diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 077e6f437d..fb8bd5ccd4 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -146,6 +146,7 @@ def _generate_start_stencil(self) -> None: acc_present=stencil.acc_present, mergecopy=stencil.mergecopy, copies=stencil.copies, + optional_module=stencil.optional_module, ) i += 2 diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index 93903cab78..07cb43872e 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -89,6 +89,7 @@ class BaseStartStencilData(CodeGenInput): class StartStencilData(BaseStartStencilData): mergecopy: Optional[bool] copies: Optional[bool] + optional_module: Optional[str] @dataclass diff --git a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py index 6860b455f5..f550b4fd09 100644 --- a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py @@ -40,6 +40,7 @@ "accpresent", "mergecopy", "copies", + "optional_module", "horizontal_lower", "horizontal_upper", "vertical_lower", diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index e6a0bd8ef8..ed1c8ad93b 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -10,9 +10,10 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Any +from typing import Any, Optional from icon4pytools.common.logger import setup_logger +from icon4pytools.liskov.codegen.integration.deserialise import DEFAULT_STARTSTENCIL_OPTIONAL_MODULE from icon4pytools.liskov.codegen.integration.interface import ( EndDeleteData, EndFusedStencilData, @@ -30,7 +31,18 @@ logger = setup_logger(__name__) -class StencilTransformer(Step): +def _remove_stencils( + parsed: IntegrationCodeInterface, stencils_to_remove: list[CodeGenInput] +) -> None: + attributes_to_modify = ["StartStencil", "EndStencil"] + + for attr_name in attributes_to_modify: + current_stencil_list = getattr(parsed, attr_name) + modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] + setattr(parsed, attr_name, modified_stencil_list) + + +class FusedStencilTransformer(Step): def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: self.parsed = parsed self.fused = fused @@ -71,7 +83,7 @@ def _process_stencils_for_deletion(self) -> None: self._create_delete_directives(start_single, end_single) stencils_to_remove += [start_single, end_single] - self._remove_stencils(stencils_to_remove) + _remove_stencils(self.parsed, stencils_to_remove) def _stencil_is_removable( self, @@ -105,14 +117,6 @@ def _create_delete_directives( directive.append(cls(startln=param.startln)) setattr(self.parsed, attr, directive) - def _remove_stencils(self, stencils_to_remove: list[CodeGenInput]) -> None: - attributes_to_modify = ["StartStencil", "EndStencil"] - - for attr_name in attributes_to_modify: - current_stencil_list = getattr(self.parsed, attr_name) - modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] - setattr(self.parsed, attr_name, modified_stencil_list) - def _remove_fused_stencils(self) -> None: self.parsed.StartFusedStencil = [] self.parsed.EndFusedStencil = [] @@ -120,3 +124,52 @@ def _remove_fused_stencils(self) -> None: def _remove_delete(self) -> None: self.parsed.StartDelete = [] self.parsed.EndDelete = [] + + +class OptionalModulesTransformer(Step): + def __init__( + self, parsed: IntegrationCodeInterface, optional_modules_to_enable: Optional[tuple[str]] + ) -> None: + self.parsed = parsed + self.optional_modules_to_enable = optional_modules_to_enable + + def __call__(self, data: Any = None) -> IntegrationCodeInterface: + """Transform stencils in the parse tree based on 'optional_modules_to_enable', either enabling specific modules or removing them. + + Args: + data (Any): Optional data to be passed. Defaults to None. + + Returns: + IntegrationCodeInterface: The modified interface object. + """ + if self.optional_modules_to_enable is not None: + action = "enabling" + else: + action = "removing" + logger.info(f"Transforming stencils by {action} optional modules.") + self._transform_stencils() + + return self.parsed + + def _transform_stencils(self) -> None: + """Identify stencils to transform based on 'optional_modules_to_enable' and applies necessary changes.""" + stencils_to_remove = [] + for start_stencil, end_stencil in zip( + self.parsed.StartStencil, self.parsed.EndStencil, strict=False + ): + if self._should_remove_stencil(start_stencil): + stencils_to_remove.extend([start_stencil, end_stencil]) + + _remove_stencils(self.parsed, stencils_to_remove) + + def _should_remove_stencil(self, stencil: StartStencilData) -> bool: + """Determine if a stencil should be removed based on 'optional_modules_to_enable'. + + Returns: + bool: True if the stencil should be removed, False otherwise. + """ + if stencil.optional_module == DEFAULT_STARTSTENCIL_OPTIONAL_MODULE: + return False + if self.optional_modules_to_enable is None: + return True + return stencil.optional_module not in self.optional_modules_to_enable diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 4d458ad1cc..8b924926b3 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path -from typing import Any +from typing import Any, Optional from icon4pytools.liskov.codegen.integration.deserialise import IntegrationCodeDeserialiser from icon4pytools.liskov.codegen.integration.generate import IntegrationCodeGenerator @@ -22,7 +22,10 @@ from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser from icon4pytools.liskov.parsing.scan import DirectivesScanner -from icon4pytools.liskov.parsing.transform import StencilTransformer +from icon4pytools.liskov.parsing.transform import ( + FusedStencilTransformer, + OptionalModulesTransformer, +) from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline @@ -70,21 +73,28 @@ def parse_fortran_file( @linear_pipeline -def process_stencils(parsed: IntegrationCodeInterface, fused: bool) -> list[Step]: +def process_stencils( + parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: Optional[tuple[str]] +) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. This function takes an input `parsed` object of type `IntegrationCodeInterface` and a `fused` boolean flag. - It then executes a linear pipeline, consisting of two steps: transformation of stencils for fusion or unfusion, - and updating fields with information from GT4Py stencils. + It then executes a linear pipeline, consisting of three steps: transformation of stencils for fusion or unfusion, + enabling optional modules, and updating fields with information from GT4Py stencils. Args: parsed (IntegrationCodeInterface): The input object containing parsed integration code. fused (bool): A boolean flag indicating whether to produce fused (True) or unfused (False) execution. + optional_modules_to_enable (Optional[tuple[str]]): A tuple of optional modules to enable. Returns: - The updated and transformed object with fields containing information from GT4Py stencils. + The updated and transformed IntegrationCodeInterface object. """ - return [StencilTransformer(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] + return [ + FusedStencilTransformer(parsed, fused), + OptionalModulesTransformer(parsed, optional_modules_to_enable), + UpdateFieldsWithGt4PyStencils(parsed), + ] @linear_pipeline diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 8c6868a76f..18e803822a 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -14,6 +14,7 @@ import os import pkgutil import re +import traceback from importlib import reload import icon4py.model.atmosphere.diffusion.stencils as diffusion @@ -137,7 +138,9 @@ def test_codegen(cli, stencil_module, stencil_name, flags) -> None: with cli.isolated_filesystem(): cli_args = [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH, *flags] result = cli.invoke(main, cli_args) - assert result.exit_code == 0 + assert ( + result.exit_code == 0 + ), f"Codegen failed with error:\n{''.join(traceback.format_exception(*result.exc_info))}" check_code_was_generated(stencil_name) diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 55f2e5f3b2..454d01e428 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -77,6 +77,7 @@ def test_stencil_collector_invalid_member(): acc_present=False, mergecopy=False, copies=True, + optional_module="None", ) ], Imports=None, diff --git a/tools/tests/liskov/test_generation.py b/tools/tests/liskov/test_generation.py index 46b7f43dae..b715f6b643 100644 --- a/tools/tests/liskov/test_generation.py +++ b/tools/tests/liskov/test_generation.py @@ -79,6 +79,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module="None", ) end_stencil_data = EndStencilData( name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False diff --git a/tools/tests/liskov/test_serialisation_deserialiser.py b/tools/tests/liskov/test_serialisation_deserialiser.py index 0431086beb..518029af3a 100644 --- a/tools/tests/liskov/test_serialisation_deserialiser.py +++ b/tools/tests/liskov/test_serialisation_deserialiser.py @@ -105,6 +105,7 @@ def parsed_dict(): "horizontal_lower": "i_startidx", "horizontal_upper": "i_endidx", "accpresent": "True", + "optional_module": "advection", } ], "StartProfile": [{"name": "apply_nabla2_to_vn_in_lateral_boundary"}], @@ -127,5 +128,7 @@ def test_savepoint_data_factory(parsed_dict): savepoints = SavepointDataFactory()(parsed_dict) assert len(savepoints) == 2 assert any([isinstance(sp, SavepointData) for sp in savepoints]) + # check that unnecessary keys have been removed + assert not any(f.variable == "optional_module" for sp in savepoints for f in sp.fields) assert any([isinstance(f, FieldSerialisationData) for f in savepoints[0].fields]) assert any([isinstance(m, Metadata) for m in savepoints[0].metadata]) diff --git a/tools/tests/liskov/test_transform.py b/tools/tests/liskov/test_transform.py index 4c2a60454c..1ac5a2bcbb 100644 --- a/tools/tests/liskov/test_transform.py +++ b/tools/tests/liskov/test_transform.py @@ -33,7 +33,10 @@ StartProfileData, StartStencilData, ) -from icon4pytools.liskov.parsing.transform import StencilTransformer +from icon4pytools.liskov.parsing.transform import ( + FusedStencilTransformer, + OptionalModulesTransformer, +) @pytest.fixture @@ -94,6 +97,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module="None", ) end_stencil_data1 = EndStencilData( name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False @@ -126,6 +130,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module="advection", ) end_stencil_data2 = EndStencilData( name="stencil2", startln=6, noendif=False, noprofile=False, noaccenddata=False @@ -165,20 +170,32 @@ def integration_code_interface(): @pytest.fixture -def stencil_transform_fused(integration_code_interface): - return StencilTransformer(integration_code_interface, fused=True) +def fused_stencil_transform_fused(integration_code_interface): + return FusedStencilTransformer(integration_code_interface, fused=True) @pytest.fixture -def stencil_transform_unfused(integration_code_interface): - return StencilTransformer(integration_code_interface, fused=False) +def fused_stencil_transform_unfused(integration_code_interface): + return FusedStencilTransformer(integration_code_interface, fused=False) + + +@pytest.fixture +def optional_modules_transform_enabled(integration_code_interface): + return OptionalModulesTransformer( + integration_code_interface, optional_modules_to_enable=["advection"] + ) + + +@pytest.fixture +def optional_modules_transform_disabled(integration_code_interface): + return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable=None) def test_transform_fused( - stencil_transform_fused, + fused_stencil_transform_fused, ): # Check that the transformed interface is as expected - transformed = stencil_transform_fused() + transformed = fused_stencil_transform_fused() assert len(transformed.StartFusedStencil) == 1 assert len(transformed.EndFusedStencil) == 1 assert len(transformed.StartStencil) == 1 @@ -188,10 +205,10 @@ def test_transform_fused( def test_transform_unfused( - stencil_transform_unfused, + fused_stencil_transform_unfused, ): # Check that the transformed interface is as expected - transformed = stencil_transform_unfused() + transformed = fused_stencil_transform_unfused() assert not transformed.StartFusedStencil assert not transformed.EndFusedStencil @@ -199,3 +216,30 @@ def test_transform_unfused( assert len(transformed.EndStencil) == 2 assert not transformed.StartDelete assert not transformed.EndDelete + + +def test_transform_optional_enabled( + optional_modules_transform_enabled, +): + # Check that the transformed interface is as expected + transformed = optional_modules_transform_enabled() + assert len(transformed.StartFusedStencil) == 1 + assert len(transformed.EndFusedStencil) == 1 + assert len(transformed.StartStencil) == 2 + assert len(transformed.EndStencil) == 2 + assert len(transformed.StartDelete) == 1 + assert len(transformed.EndDelete) == 1 + + +def test_transform_optional_disabled( + optional_modules_transform_disabled, +): + # Check that the transformed interface is as expected + transformed = optional_modules_transform_disabled() + + assert len(transformed.StartFusedStencil) == 1 + assert len(transformed.EndFusedStencil) == 1 + assert len(transformed.StartStencil) == 1 + assert len(transformed.EndStencil) == 1 + assert len(transformed.StartDelete) == 1 + assert len(transformed.EndDelete) == 1