Porting the tracer advection to gt4py (#252)

Merge advection to main

.github/workflows/icon4py-qa.yml

@@ -58,4 +58,6 @@ jobs:
             - name: Run checks icon4py-model-atmosphere-diffusion
               run: |
                   pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files
-
+            - name: Run checks icon4py-model-atmosphere-advection
+              run: |
+                  pre-commit run --config model/atmosphere/advection/.pre-commit-config.yaml --all-files

model/README.md

@@ -6,6 +6,7 @@ It includes the following packages:
 
 - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model
 - `atmosphere/diffusion`: Contains the implementation of diffusion in the ICON model
+- `atmosphere/advection`: Contains implementations of the advection component of the ICON model
 - `common`: Contains shared functionality that is required by multiple components.
 - `driver`: Contains the driving code for the model
 
@@ -18,6 +19,7 @@ In the following example it is assumed that you have already created and activat
 ```bash
 # changing into the corresponding directory
 cd model/atmosphere/dycore
+cd model/atmosphere/advection
 
 # installing a development version
 pip install -r requirements-dev.txt

model/atmosphere/advection/.flake8

@@ -0,0 +1,42 @@
+[flake8]
+# Some sane defaults for the code style checker flake8
+max-line-length = 100
+max-complexity = 15
+doctests = true
+extend-ignore =
+    # Do not perform function calls in argument defaults
+    B008,
+    # Public code object needs docstring
+    D1,
+    # Disable dargling errors by default
+    DAR,
+    # Whitespace before ':' (black formatter breaks this sometimes)
+    E203,
+    # Line too long (using Bugbear's B950 warning)
+    E501,
+    # Line break occurred before a binary operator
+    W503
+
+exclude =
+    .eggs,
+    .gt_cache,
+    .ipynb_checkpoints,
+    .tox,
+    _local_,
+    build,
+    dist,
+    docs,
+    _external_src,
+    tests/_disabled,
+    setup.py
+
+rst-roles =
+    py:mod, mod,
+    py:func, func,
+    py:data, data,
+    py:const, const,
+    py:class, class,
+    py:meth, meth,
+    py:attr, attr,
+    py:exc, exc,
+    py:obj, obj,

model/atmosphere/advection/.pre-commit-config.yaml

@@ -0,0 +1,114 @@
+# NOTE: pre-commit runs all hooks from the root folder of the repository,
+# as regular git hooks do. Therefore, paths passed as arguments to the plugins
+# should always be relative to the root folder.
+
+default_stages: [commit, push]
+default_language_version:
+  python: python3.10
+minimum_pre_commit_version: 2.20.0
+files: "model/atmosphere/advection/.*"
+
+repos:
+- repo: meta
+  hooks:
+  - id: check-hooks-apply
+    stages: [manual]
+  - id: check-useless-excludes
+    stages: [manual]
+
+- repo: https://github.com/asottile/setup-cfg-fmt
+  rev: v1.20.1
+  hooks:
+          # Run only manually because it deletes comments
+  - id: setup-cfg-fmt
+    name: format setup.cfg
+    stages: [manual]
+
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.3.0
+  hooks:
+  - id: check-case-conflict
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-symlinks
+  - id: check-yaml
+  - id: debug-statements
+  - id: destroyed-symlinks
+
+- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
+  rev: v2.6.0
+  hooks:
+  - id: pretty-format-ini
+    args: [--autofix]
+  - id: pretty-format-toml
+    args: [--autofix]
+  - id: pretty-format-yaml
+    args: [--autofix, --preserve-quotes, --indent, "2"]
+
+- repo: https://github.com/pre-commit/mirrors-prettier
+  rev: v3.0.0-alpha.4
+  hooks:
+  - id: prettier
+    types_or: [markdown, json]
+
+- repo: https://github.com/Lucas-C/pre-commit-hooks
+  rev: v1.3.0
+  hooks:
+  - id: insert-license
+    name: add license for all ICON4Py Python source files
+    types: [python]
+    args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo]
+
+- repo: https://github.com/asottile/yesqa
+  rev: v1.3.0
+  hooks:
+  - id: yesqa
+
+- repo: https://github.com/psf/black
+  rev: '22.3.0'
+  hooks:
+  - id: black
+    name: black Python formatter
+    args: [--config, model/atmosphere/advection/pyproject.toml]
+
+- repo: https://github.com/asottile/blacken-docs
+  rev: v1.12.1
+  hooks:
+  - id: blacken-docs
+    name: black Python formatter for docstrings
+    additional_dependencies: [black==22.3.0]
+
+- repo: https://github.com/PyCQA/isort
+  rev: '5.12.0'
+  hooks:
+  - id: isort
+    args: [--config-root, model/atmosphere/advection/, --resolve-all-configs]
+
+- repo: https://github.com/PyCQA/flake8
+  rev: '4.0.1'
+  hooks:
+  - id: flake8
+    name: flake8 code style checks
+    additional_dependencies:
+    - darglint
+    - flake8-bugbear
+    - flake8-builtins
+    - flake8-debugger
+    - flake8-docstrings
+    - flake8-eradicate
+    - flake8-mutable
+    - pygments
+    args: [--config=model/atmosphere/advection/.flake8, model/atmosphere/advection/src/icon4py/]
+
+- repo: local
+  hooks:
+  - id: mypy
+    name: mypy static type checker
+    entry: bash -c 'echo mypy temporarily disabled'
+            #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' --
+    language: system
+    types_or: [python, pyi]
+    always_run: true
+            #pass_filenames: false
+    require_serial: true
+    stages: [commit]

model/atmosphere/advection/README.md

@@ -0,0 +1,9 @@
+# icon4py-atmosphere-advection
+
+## Description
+
+Contains code ported from ICON `src/advection`, which is the advection component of the ICON model.
+
+## Installation instructions
+
+Check the `README.md` at the root of the `model` folder for installation instructions.

model/atmosphere/advection/advection_tests/__init__.py

@@ -0,0 +1,12 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later

model/atmosphere/advection/advection_tests/conftest.py

@@ -0,0 +1,33 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+import pytest
+from gt4py.next.program_processors.runners.roundtrip import executor
+
+from icon4py.model.common.test_utils.simple_mesh import SimpleMesh
+
+
+BACKENDS = {"embedded": executor}
+MESHES = {"simple_mesh": SimpleMesh()}
+
+
+@pytest.fixture(
+    ids=MESHES.keys(),
+    params=MESHES.values(),
+)
+def mesh(request):
+    return request.param
+
+
+@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values())
+def backend(request):
+    return request.param

model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py

@@ -0,0 +1,63 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+
+from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01
+from icon4py.model.common.dimension import EdgeDim, KDim
+from icon4py.model.common.test_utils.helpers import random_field, zero_field
+from icon4py.model.common.test_utils.simple_mesh import SimpleMesh
+
+
+def btraj_dreg_stencil_01_numpy(
+    lcounterclock: bool,
+    p_vn: np.array,
+    tangent_orientation: np.array,
+):
+    tangent_orientation = np.expand_dims(tangent_orientation, axis=-1)
+
+    tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape)
+
+    lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False)
+
+    mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape)
+
+    lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False)
+
+    return lvn_sys_pos
+
+
+def test_btraj_dreg_stencil_01():
+    mesh = SimpleMesh()
+    lcounterclock = True
+    p_vn = random_field(mesh, EdgeDim, KDim)
+
+    tangent_orientation = random_field(mesh, EdgeDim)
+
+    lvn_sys_pos = zero_field(mesh, EdgeDim, KDim, dtype=bool)
+
+    ref = btraj_dreg_stencil_01_numpy(
+        lcounterclock,
+        np.asarray(p_vn),
+        np.asarray(tangent_orientation),
+    )
+
+    btraj_dreg_stencil_01(
+        lcounterclock,
+        p_vn,
+        tangent_orientation,
+        lvn_sys_pos,
+        offset_provider={},
+    )
+
+    assert np.allclose(ref, lvn_sys_pos)

model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py

@@ -0,0 +1,71 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+from gt4py.next.ffront.fbuiltins import int32
+from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider
+
+from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02
+from icon4py.model.common.dimension import ECDim, EdgeDim, KDim
+from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field
+from icon4py.model.common.test_utils.simple_mesh import SimpleMesh
+
+
+def btraj_dreg_stencil_02_numpy(
+    p_vn: np.array,
+    p_vt: np.array,
+    edge_cell_length: np.array,
+    p_dt: float,
+):
+    lvn_pos = np.where(p_vn >= 0.0, True, False)
+
+    traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt
+
+    edge_cell_length = np.expand_dims(edge_cell_length, axis=-1)
+    e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1])
+
+    opt_famask_dsl = np.where(
+        traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)),
+        int32(1),
+        int32(0),
+    )
+
+    return opt_famask_dsl
+
+
+def test_btraj_dreg_stencil_02():
+    mesh = SimpleMesh()
+    p_vn = random_field(mesh, EdgeDim, KDim)
+    p_vt = random_field(mesh, EdgeDim, KDim)
+    edge_cell_length = np.asarray(mesh.e2c, dtype=float)
+    edge_cell_length_new = as_1D_sparse_field(edge_cell_length, ECDim)
+    p_dt = 1.0
+    opt_famask_dsl = zero_field(mesh, EdgeDim, KDim, dtype=int32)
+
+    ref = btraj_dreg_stencil_02_numpy(
+        np.asarray(p_vn), np.asarray(p_vt), np.asarray(edge_cell_length), p_dt
+    )
+
+    btraj_dreg_stencil_02(
+        p_vn,
+        p_vt,
+        edge_cell_length_new,
+        p_dt,
+        opt_famask_dsl,
+        offset_provider={
+            "E2C": mesh.get_e2c_offset_provider(),
+            "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c),
+        },
+    )
+
+    assert np.allclose(ref, opt_famask_dsl)