diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 18be464b9..55b260cad 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -21,10 +21,14 @@ on: types: [submitted] jobs: - pre-commit-icon4py-model-common: + pre-commit-icon4py-model: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: @@ -36,8 +40,8 @@ jobs: **/base-requirements-dev.txt **/requirements.txt **/requirements-dev.txt - - name: Install icon4py-common - working-directory: model/common + - name: Install icon4py-model packages + working-directory: model run: | python -m pip install --upgrade pip setuptools wheel python -m pip install -r ./requirements-dev.txt @@ -45,27 +49,13 @@ jobs: - name: Run checks icon4py-model-common run: | pre-commit run --config model/common/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-atmosphere-dycore: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-atmosphere-dycore - working-directory: model/atmosphere/dycore + - name: Run checks icon4py-model-driver run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt + pre-commit run --config model/driver/.pre-commit-config.yaml --all-files - name: Run checks icon4py-model-atmosphere-dycore run: | pre-commit run --config model/atmosphere/dycore/.pre-commit-config.yaml --all-files + - name: Run checks icon4py-model-atmosphere-diffusion + run: | + pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files + diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index 0f550ce84..af1138fed 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -31,6 +31,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install boost + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: diff --git a/.gitignore b/.gitignore index 56ff4f8b6..e7620d5cf 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,8 @@ _local _external_src _reports tmp +testdata +simple_mesh*.nc ### GT4Py #### .gt_cache/ diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 3c2f34e7a..991b3d473 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,7 @@ # VCS -e git+https://github.com/GridTools/gt4py.git@main#egg=gt4py +git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python +git+https://github.com/ghex-org/GHEX.git#subdirectory=bindings/python # PyPI bump2version>=1.0.1 @@ -13,6 +15,7 @@ flake8-eradicate>=1.3.0 flake8-mutable>=1.2.0 isort~=5.10 mypy>=0.942 +typing-extensions==4.5.0 pre-commit~=2.15 pytest>=6.1 pytest-benchmark>=4.0.0 @@ -20,6 +23,8 @@ pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 pytest-xdist[psutil]>=2.2 +pytest-mpi>=0.6 setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 +wget>=3.2 diff --git a/base-requirements.txt b/base-requirements.txt index 4e204ea2f..0720ba258 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,2 +1,3 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main +pyghex @ git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python diff --git a/model/README.md b/model/README.md index d59191505..0f076a932 100644 --- a/model/README.md +++ b/model/README.md @@ -5,7 +5,9 @@ This folder contains Python implementations for multiple ICON components. It includes the following packages: - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model +- `atmosphere/diffusion`: Contains the implementation of diffusion in the ICON model - `common`: Contains shared functionality that is required by multiple components. +- `driver`: Contains the driving code for the model ## Installation Instructions @@ -22,3 +24,47 @@ pip install -r requirements-dev.txt ``` **Note**: For more information specific to each component, please refer to the README in their respective subfolders. + +### Testing + +See the repository [README](../README.md) for general information. + +#### Data dependent tests + +Some test depend on serialized data generated by a full model run. +Those test are marked with `pytest.mark.datatest` and are only run when the `--datatest` +option is specified. Note that due to `pytests` configuration discovery +you need to specify the base a package directory i.e. one that contains a `pyproject.toml` for the +commandline option to work. + +```bash + +pytest -v --datatest model/common +pytest -v --datatest model/atmosphere/diffusion +``` + +#### Testing parallel code + +Tests for parallel codes using MPI are collected in specific subpackages of the model components test folders (e.g. `diffusion_tests/mpi_tests`). All parallel tests are marked with `@pytest.mark.mpi` and are skipped if the `--with-mpi` is not passed option is not passed to `pytest` In order to run them, you need a MPI installation on your system: On Debian-Linux do + +```bash +sudo apt-get install libopenmpi-dev +``` + +or + +```bash +sudo apt-get install mpich +``` + +on MacOs + +```bash +brew install mpich +``` + +Then you can run the tests with + +```bash +mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests +``` diff --git a/model/atmosphere/diffusion/.bumpversion.cfg b/model/atmosphere/diffusion/.bumpversion.cfg new file mode 100644 index 000000000..ef8bea214 --- /dev/null +++ b/model/atmosphere/diffusion/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/atmosphere/diffusion/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/atmosphere/diffusion/.flake8 b/model/atmosphere/diffusion/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/atmosphere/diffusion/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/atmosphere/diffusion/.pre-commit-config.yaml b/model/atmosphere/diffusion/.pre-commit-config.yaml new file mode 100644 index 000000000..1d6ce2a03 --- /dev/null +++ b/model/atmosphere/diffusion/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/atmosphere/diffusion/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/atmosphere/diffusion/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/atmosphere/diffusion/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/atmosphere/diffusion/.flake8, model/atmosphere/diffusion/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/atmosphere/diffusion/README.md b/model/atmosphere/diffusion/README.md new file mode 100644 index 000000000..51a415319 --- /dev/null +++ b/model/atmosphere/diffusion/README.md @@ -0,0 +1,9 @@ +# icon4py-atmosphere-diffusion + +## Description + +Python port of ICON diffusion module. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/atmosphere/diffusion/diffusion_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/diffusion/diffusion_tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py new file mode 100644 index 000000000..d42920444 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/conftest.py @@ -0,0 +1,94 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 # import fixtures from test_utils package + backend, + damping_height, + data_provider, + datapath, + decomposition_info, + download_ser_data, + grid_savepoint, + icon_grid, + interpolation_savepoint, + linit, + mesh, + metrics_savepoint, + ndyn_substeps, + processor_props, + ranked_data_path, + step_date_exit, + step_date_init, +) + + +@pytest.fixture +def r04b09_diffusion_config( + ndyn_substeps, # noqa: F811 # imported `ndyn_substeps` fxiture +) -> DiffusionConfig: + """ + Create DiffusionConfig matching MCH_CH_r04b09_dsl. + + Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + from the default. + """ + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + n_substeps=ndyn_substeps, + ) + + +@pytest.fixture +def diffusion_savepoint_init( + data_provider, # noqa: F811 # imported fixtures data_provider + linit, # noqa: F811 # imported fixtures linit + step_date_init, # noqa: F811 # imported fixtures data_provider +): + """ + Load data from ICON savepoint at start of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' + fixture, passing 'step_date_init=' + + linit flag can be set by overriding the 'linit' fixture + """ + return data_provider.from_savepoint_diffusion_init(linit=linit, date=step_date_init) + + +@pytest.fixture +def diffusion_savepoint_exit( + data_provider, # noqa: F811 # imported fixtures data_provider` + linit, # noqa: F811 # imported fixtures linit` + step_date_exit, # noqa: F811 # imported fixtures step_date_exit` +): + """ + Load data from ICON savepoint at exist of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_data' + fixture, passing 'step_data=' + """ + sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) + return sp diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py new file mode 100644 index 000000000..8a6602799 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -0,0 +1,110 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size + +from ..utils import verify_diffusion_fields + + +@pytest.mark.mpi +@pytest.mark.parametrize("ndyn_substeps", [2]) +@pytest.mark.parametrize("linit", [True, False]) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_parallel_diffusion( + r04b09_diffusion_config, + step_date_init, + linit, + ndyn_substeps, + processor_props, + decomposition_info, + icon_grid, + diffusion_savepoint_init, + diffusion_savepoint_exit, + grid_savepoint, + metrics_savepoint, + interpolation_savepoint, + damping_height, +): + check_comm_size(processor_props) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: decomposition info : klevels = {decomposition_info.klevels}, " + f"local cells = {decomposition_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " + f"local vertices = {decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + ) + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + + diffusion_params = DiffusionParams(r04b09_diffusion_config) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + exchange = create_exchange(processor_props, decomposition_info) + + diffusion = Diffusion(exchange) + + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=diffusion_params, + vertical_params=VerticalModelParams(grid_savepoint.vct_a(), damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + if linit: + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + else: + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + + verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" + ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py similarity index 95% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index d9d058adf..02e618fcd 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( apply_nabla2_and_nabla4_global_to_vn, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py similarity index 62% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py index 4b4a9cecc..29d3fa391 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -14,7 +14,10 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( + apply_nabla2_and_nabla4_global_to_vn, +) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( apply_nabla2_and_nabla4_to_vn, ) from icon4py.model.common.dimension import EdgeDim, KDim @@ -67,3 +70,34 @@ def reference( - diff_multfac_vn * z_nabla4_e2 * area_edge ) return dict(vn=vn) + + +class TestApplyNabla2AndNabla4ToVnGlobalMode(StencilTest): + PROGRAM = apply_nabla2_and_nabla4_global_to_vn + + OUTPUTS = ("vn",) + + @pytest.fixture + def input_data(self, mesh): + area_edge = random_field(mesh, EdgeDim) + kh_smag_e = random_field(mesh, EdgeDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) + diff_multfac_vn = random_field(mesh, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + return dict( + area_edge=area_edge, + kh_smag_e=kh_smag_e, + z_nabla2_e=z_nabla2_e, + z_nabla4_e2=z_nabla4_e2, + diff_multfac_vn=diff_multfac_vn, + vn=vn, + ) + + @staticmethod + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py similarity index 94% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index 792882432..ff44a1f67 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( apply_nabla2_to_vn_in_lateral_boundary, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py similarity index 80% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py index c3f472521..147fd5c59 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -30,7 +31,8 @@ def reference( z_nabla2_c: np.array, geofac_n2s: np.array, w: np.array, - diff_multfac_w, + diff_multfac_w: float, + **kwargs, ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) @@ -43,12 +45,14 @@ def input_data(self, mesh): z_nabla2_c = random_field(mesh, CellDim, KDim) geofac_n2s = random_field(mesh, CellDim, C2E2CODim) w = random_field(mesh, CellDim, KDim) - diff_multfac_w = 5.0 - return dict( area=area, z_nabla2_c=z_nabla2_c, geofac_n2s=geofac_n2s, w=w, - diff_multfac_w=diff_multfac_w, + diff_multfac_w=5.0, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py similarity index 83% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 98193e058..6608b9008 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, ) from icon4py.model.common.dimension import CellDim, KDim @@ -37,6 +38,10 @@ def input_data(self, mesh): diff_multfac_n2w=diff_multfac_n2w, cell_area=cell_area, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) @staticmethod @@ -46,6 +51,7 @@ def reference( diff_multfac_n2w: np.array, cell_area: np.array, z_nabla2_c: np.array, + **kwargs, ) -> np.array: cell_area = np.expand_dims(cell_area, axis=-1) w = w + diff_multfac_n2w * cell_area * z_nabla2_c diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py similarity index 94% rename from model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py index c62891857..c70fc8fa2 100644 --- a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( calculate_diagnostics_for_turbulence, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py similarity index 85% rename from model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py index f846a7a78..c7a7016f5 100644 --- a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -13,8 +13,9 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( calculate_horizontal_gradients_for_turbulence, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim @@ -50,4 +51,8 @@ def input_data(self, mesh): geofac_grg_y=geofac_grg_y, dwdx=dwdx, dwdy=dwdy, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 98% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 184dba0dc..b230948b5 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_and_smag_coefficients_for_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py similarity index 81% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py index 45e4327a9..310a91c21 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py @@ -13,8 +13,11 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import calculate_nabla2_for_w +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( + calculate_nabla2_for_w, +) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -39,4 +42,8 @@ def input_data(self, mesh): w=w, geofac_n2s=geofac_n2s, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py similarity index 94% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py index 7772dd5db..a82916ce3 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py @@ -14,7 +14,9 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import calculate_nabla2_for_z +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + calculate_nabla2_for_z, +) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py similarity index 93% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py index 5b1f99875..0d6bd5251 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py @@ -14,7 +14,9 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import calculate_nabla2_of_theta +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + calculate_nabla2_of_theta, +) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( StencilTest, diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py similarity index 97% rename from model/atmosphere/dycore/tests/test_calculate_nabla4.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py index 07d097bf1..de26d3cd6 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla4 import calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.model.common.test_utils.helpers import ( StencilTest, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py new file mode 100644 index 000000000..383d10f34 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -0,0 +1,348 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionInitSavepoint + +from .utils import ( + diff_multfac_vn_numpy, + enhanced_smagorinski_factor_numpy, + smag_limit_numpy, + verify_diffusion_fields, +) + + +def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.hdiff_efdt_ratio = 1.0 + config.hdiff_w_efdt_ratio = 2.0 + + params = DiffusionParams(config) + + assert params.K2 == pytest.approx(0.125, abs=1e-12) + assert params.K4 == pytest.approx(0.125 / 8.0, abs=1e-12) + assert params.K6 == pytest.approx(0.125 / 64.0, abs=1e-12) + assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) + + +def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.hdiff_efdt_ratio = 0.0 + config.hdiff_w_efdt_ratio = 0.0 + + params = DiffusionParams(config) + + assert params.K2 == 0.0 + assert params.K4 == 0.0 + assert params.K6 == 0.0 + assert params.K4W == 0.0 + + +def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): + config = r04b09_diffusion_config + config.smagorinski_scaling_factor = 0.15 + config.diffusion_type = 4 + + params = DiffusionParams(config) + assert len(params.smagorinski_factor) == 1 + assert params.smagorinski_factor[0] == pytest.approx(0.15, abs=1e-16) + assert params.smagorinski_height is None + + +def test_smagorinski_heights_diffusion_type_5_are_consistent( + r04b09_diffusion_config, +): + config = r04b09_diffusion_config + config.smagorinski_scaling_factor = 0.15 + config.diffusion_type = 5 + + params = DiffusionParams(config) + assert len(params.smagorinski_height) == 4 + assert min(params.smagorinski_height) == params.smagorinski_height[0] + assert max(params.smagorinski_height) == params.smagorinski_height[-1] + assert params.smagorinski_height[0] < params.smagorinski_height[1] + assert params.smagorinski_height[1] < params.smagorinski_height[3] + assert params.smagorinski_height[2] != params.smagorinski_height[1] + assert params.smagorinski_height[2] != params.smagorinski_height[3] + + +def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): + params = DiffusionParams(r04b09_diffusion_config) + assert len(params.smagorinski_factor) == len(params.smagorinski_height) + assert len(params.smagorinski_factor) == 4 + assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) + + +@pytest.mark.datatest +def test_diffusion_init( + diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + step_date_init, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + + meta = diffusion_savepoint_init.get_metadata("linit", "date") + + assert meta["linit"] is False + assert meta["date"] == step_date_init + + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + assert diffusion.diff_multfac_w == min( + 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float + ) + + assert np.allclose(0.0, np.asarray(diffusion.v_vert)) + assert np.allclose(0.0, np.asarray(diffusion.u_vert)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) + + shape_k = (icon_grid.n_lev(),) + expected_smag_limit = smag_limit_numpy( + diff_multfac_vn_numpy, + shape_k, + additional_parameters.K4, + config.substep_as_float, + ) + + assert diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float + assert np.allclose(expected_smag_limit, diffusion.smag_limit) + + expected_diff_multfac_vn = diff_multfac_vn_numpy( + shape_k, additional_parameters.K4, config.substep_as_float + ) + assert np.allclose(expected_diff_multfac_vn, diffusion.diff_multfac_vn) + expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( + additional_parameters.smagorinski_factor, + additional_parameters.smagorinski_height, + grid_savepoint.vct_a(), + ) + assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) + + +def _verify_init_values_against_savepoint( + savepoint: IconDiffusionInitSavepoint, diffusion: Diffusion +): + dtime = savepoint.get_metadata("dtime")["dtime"] + + assert savepoint.nudgezone_diff() == diffusion.nudgezone_diff + assert savepoint.bdy_diff() == diffusion.bdy_diff + assert savepoint.fac_bdydiff_v() == diffusion.fac_bdydiff_v + assert savepoint.smag_offset() == diffusion.smag_offset + assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w + + # this is done in diffusion.run(...) because it depends on the dtime + scale_k(diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={}) + assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) + + assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) + assert np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) + + +@pytest.mark.datatest +def test_verify_diffusion_init_against_first_regular_savepoint( + diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + r04b09_diffusion_config, + icon_grid, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + vct_a = grid_savepoint.vct_a() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() + + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=VerticalModelParams(vct_a, damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) + + +@pytest.mark.datatest +@pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) +def test_verify_diffusion_init_against_other_regular_savepoint( + r04b09_diffusion_config, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + diffusion_savepoint_init, + damping_height, +): + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() + + diffusion = Diffusion() + diffusion.init( + icon_grid, + config, + additional_parameters, + vertical_params, + metric_state, + interpolation_state, + edge_params, + cell_params, + ) + + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "step_date_init, step_date_exit", + [ + ("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), + ], +) +def test_run_diffusion_single_step( + diffusion_savepoint_init, + diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, +): + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + vct_a = grid_savepoint.vct_a() + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_init) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit) + + +@pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) +def test_run_diffusion_initial_step( + diffusion_savepoint_init, + diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, +): + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + vct_a = grid_savepoint.vct_a() + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) + config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + + verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py new file mode 100644 index 000000000..243ab01e2 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + + +@pytest.mark.datatest +def test_verify_geofac_n2s_field_manipulation(interpolation_savepoint, icon_grid): + geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) + int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + geofac_c = np.asarray(int_state.geofac_n2s_c) + geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) + assert np.count_nonzero(geofac_nbh) > 0 + cec_table = icon_grid.get_c2cec_connectivity().table + assert np.allclose(geofac_c, geofac_n2s[:, 0]) + assert geofac_nbh[cec_table].shape == geofac_n2s[:, 1:].shape + assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py new file mode 100644 index 000000000..0ba7ce353 --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py @@ -0,0 +1,147 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + _en_smag_fac_for_zero_nshift, + _setup_runtime_diff_multfac_vn, + _setup_smag_limit, + scale_k, + set_zero_v_k, + setup_fields_for_initial_step, +) +from icon4py.model.common.dimension import KDim, VertexDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + +from .utils import diff_multfac_vn_numpy, enhanced_smagorinski_factor_numpy, smag_limit_numpy + + +def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): + return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) + + +def test_scale_k(): + mesh = SimpleMesh() + field = random_field(mesh, KDim) + scaled_field = zero_field(mesh, KDim) + factor = 2.0 + scale_k(field, factor, scaled_field, offset_provider={}) + assert np.allclose(factor * np.asarray(field), scaled_field) + + +def test_diff_multfac_vn_and_smag_limit_for_initial_step(): + mesh = SimpleMesh() + diff_multfac_vn_init = zero_field(mesh, KDim) + smag_limit_init = zero_field(mesh, KDim) + k4 = 1.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn_init).shape + + expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio) + expected_smag_limit_init = smag_limit_numpy( + initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio + ) + + setup_fields_for_initial_step( + k4, efdt_ratio, diff_multfac_vn_init, smag_limit_init, offset_provider={} + ) + + assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) + assert np.allclose(expected_smag_limit_init, smag_limit_init) + + +def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 1.0 + substeps = 5.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn).shape + + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + + _setup_runtime_diff_multfac_vn(k4, efdt_ratio, out=diff_multfac_vn, offset_provider={}) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 0.003 + substeps = 1.0 + + shape = np.asarray(diff_multfac_vn).shape + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + _setup_runtime_diff_multfac_vn(k4, substeps, out=diff_multfac_vn, offset_provider={}) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def test_init_enh_smag_fac(): + mesh = SimpleMesh() + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0, extend={KDim: 1}) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + + enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) + + _en_smag_fac_for_zero_nshift(a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim}) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) + + +def test_set_zero_vertex_k(): + mesh = SimpleMesh() + f = random_field(mesh, VertexDim, KDim) + set_zero_v_k(f, offset_provider={}) + assert np.allclose(0.0, f) + + +@pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) +def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, linit +): + savepoint = diffusion_savepoint_init + config = r04b09_diffusion_config + + params = DiffusionParams(config) + expected_diff_multfac_vn = savepoint.diff_multfac_vn() + expected_smag_limit = savepoint.smag_limit() + exptected_smag_offset = savepoint.smag_offset() + + diff_multfac_vn = zero_field(icon_grid, KDim) + smag_limit = zero_field(icon_grid, KDim) + setup_fields_for_initial_step( + params.K4, + config.hdiff_efdt_ratio, + diff_multfac_vn, + smag_limit, + offset_provider={}, + ) + assert np.allclose(expected_smag_limit, smag_limit) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert exptected_smag_offset == 0.0 diff --git a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 92% rename from model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 50e92d4f5..554142226 100644 --- a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( enhance_diffusion_coefficient_for_grid_point_cold_pools, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 94% rename from model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index 61966f1ad..84d8ed516 100644 --- a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py similarity index 74% rename from model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py index 43cd3f4aa..6b044d2fb 100644 --- a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -14,11 +14,16 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + zero_field, +) class TestTemporaryFieldsForTurbulenceDiagnostics(StencilTest): @@ -36,12 +41,11 @@ def reference( **kwargs, ) -> dict: geofac_div = np.expand_dims(geofac_div, axis=-1) - vn_geofac = vn[mesh.c2e] * geofac_div + vn_geofac = vn[mesh.c2e] * geofac_div[mesh.get_c2ce_offset_provider().table] div = np.sum(vn_geofac, axis=1) - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) diff_multfac_smag = np.expand_dims(diff_multfac_smag, axis=0) - mul = kh_smag_ec[mesh.c2e] * e_bln_c_s + mul = kh_smag_ec[mesh.c2e] * e_bln_c_s[mesh.get_c2ce_offset_provider().table] summed = np.sum(mul, axis=1) kh_c = summed / diff_multfac_smag @@ -50,9 +54,9 @@ def reference( @pytest.fixture def input_data(self, mesh): vn = random_field(mesh, EdgeDim, KDim) - geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) kh_smag_ec = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) + e_bln_c_s = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) diff_multfac_smag = random_field(mesh, KDim) kh_c = zero_field(mesh, CellDim, KDim) diff --git a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 93% rename from model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index b0114cc26..72baed394 100644 --- a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim @@ -111,6 +111,10 @@ def test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points(): vcoef_new, theta_v, z_temp, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), diff --git a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py similarity index 84% rename from model/atmosphere/dycore/tests/test_update_theta_and_exner.py rename to model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py index 729254337..871b32a89 100644 --- a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py @@ -13,8 +13,11 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.update_theta_and_exner import update_theta_and_exner +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + update_theta_and_exner, +) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -53,4 +56,8 @@ def input_data(self, mesh): theta_v=theta_v, exner=exner, rd_o_cvd=rd_o_cvd, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py new file mode 100644 index 000000000..b2fbf987d --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -0,0 +1,77 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionExitSavepoint + + +def verify_diffusion_fields( + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0 / 128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + + +def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml new file mode 100644 index 000000000..d6ab9aa52 --- /dev/null +++ b/model/atmosphere/diffusion/pyproject.toml @@ -0,0 +1,120 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] + +[project] +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "icon4py-common>=0.0.5", + "mpi4py<=3.1.4", + "pyghex>=0.3.0" +] +description = "ICON diffusion." +dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-atmosphere-diffusion" +readme = "README.md" +requires-python = ">=3.10" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'diffusion_tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +branch = true +parallel = true +source_pkgs = ['diffusion'] + +[tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +multi_line_output = 3 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +use_parentheses = true + +[tool.mypy] +disallow_incomplete_defs = true +disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] +ignore_missing_imports = false +implicit_reexport = true +install_types = true +non_interactive = true +show_column_numbers = true +show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = ['tests', 'diffusion_tests'] + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.atmosphere.diffusion.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.atmosphere.diffusion' = ['py.typed'] diff --git a/model/atmosphere/diffusion/requirements-dev.txt b/model/atmosphere/diffusion/requirements-dev.txt new file mode 100644 index 000000000..32d385e08 --- /dev/null +++ b/model/atmosphere/diffusion/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements-dev.txt +-e ../../common +-e . diff --git a/model/atmosphere/diffusion/requirements.txt b/model/atmosphere/diffusion/requirements.txt new file mode 100644 index 000000000..79787ec13 --- /dev/null +++ b/model/atmosphere/diffusion/requirements.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements.txt +../../common +. diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py new file mode 100644 index 000000000..49c96a94c --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py @@ -0,0 +1,34 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py new file mode 100644 index 000000000..d857a4f27 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -0,0 +1,850 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import functools +import logging +import math +import sys +from collections import namedtuple +from dataclasses import InitVar, dataclass, field +from enum import Enum +from typing import Final, Optional + +import numpy as np +from gt4py.next.common import Dimension +from gt4py.next.ffront.fbuiltins import Field, int32 +from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next.program_processors.runners.gtfn_cpu import ( + run_gtfn, + run_gtfn_cached, + run_gtfn_imperative, +) + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + copy_field, + init_diffusion_local_fields_for_regular_timestep, + init_nabla2_factor_in_upper_damping_zone, + scale_k, + setup_fields_for_initial_step, + zero_field, +) +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostic_quantities_for_turbulence import ( + calculate_diagnostic_quantities_for_turbulence, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( + calculate_nabla2_and_smag_coefficients_for_vn, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_theta import ( + calculate_nabla2_for_theta, +) +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points, +) +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + update_theta_and_exner, +) +from icon4py.model.common.constants import ( + CPD, + DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + GAS_CONSTANT_DRY_AIR, +) +from icon4py.model.common.decomposition.decomposed import ExchangeRuntime, SingleNode +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex +from icon4py.model.common.grid.icon_grid import IconGrid +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( + mo_intp_rbf_rbf_vec_interpol_vertex, +) + + +# flake8: noqa +log = logging.getLogger(__name__) + +VectorTuple = namedtuple("VectorTuple", "x y") + +cached_backend = run_gtfn_cached +compiled_backend = run_gtfn +imperative_backend = run_gtfn_imperative +backend = run_gtfn_cached # + + +class DiffusionType(int, Enum): + """ + Order of nabla operator for diffusion. + Note: Called `hdiff_order` in `mo_diffusion_nml.f90`. + Note: We currently only support type 5. + """ + + NO_DIFFUSION = -1 #: no diffusion + LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels + SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion + LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels + SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion + + +class DiffusionConfig: + """ + Contains necessary parameter to configure a diffusion run. + + Encapsulates namelist parameters and derived parameters. + Values should be read from configuration. + Default values are taken from the defaults in the corresponding ICON Fortran namelist files. + """ + + # TODO(Magdalena): to be read from config + # TODO(Magdalena): handle dependencies on other namelists (see below...) + + def __init__( + self, + diffusion_type: DiffusionType = DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + hdiff_temp=True, + type_vn_diffu: int = 1, + smag_3d: bool = False, + type_t_diffu: int = 2, + hdiff_efdt_ratio: float = 36.0, + hdiff_w_efdt_ratio: float = 15.0, + smagorinski_scaling_factor: float = 0.015, + n_substeps: int = 5, + zdiffu_t: bool = True, + hdiff_rcf: bool = True, + velocity_boundary_diffusion_denom: float = 200.0, + temperature_boundary_diffusion_denom: float = 135.0, + max_nudging_coeff: float = 0.02, + nudging_decay_rate: float = 2.0, + ): + """Set the diffusion configuration parameters with the ICON default values.""" + # parameters from namelist diffusion_nml + + self.diffusion_type: int = diffusion_type + + #: If True, apply diffusion on the vertical wind field + #: Called `lhdiff_w` in mo_diffusion_nml.f90 + self.apply_to_vertical_wind: bool = hdiff_w + + #: True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90 + #: Called `lhdiff_vn` in mo_diffusion_nml.f90 + self.apply_to_horizontal_wind = hdiff_vn + + #: If True, apply horizontal diffusion to temperature field + #: Called `lhdiff_temp` in mo_diffusion_nml.f90 + self.apply_to_temperature: bool = hdiff_temp + + #: If True, compute 3D Smagorinsky diffusion coefficient + #: Called `lsmag_3d` in mo_diffusion_nml.f90 + self.compute_3d_smag_coeff: bool = smag_3d + + #: Options for discretizing the Smagorinsky momentum diffusion + #: Called `itype_vn_diffu` in mo_diffusion_nml.f90 + self.type_vn_diffu: int = type_vn_diffu + + #: Options for discretizing the Smagorinsky temperature diffusion + #: Called `itype_t_diffu` inmo_diffusion_nml.f90 + self.type_t_diffu = type_t_diffu + + #: Ratio of e-folding time to (2*)time step + #: Called `hdiff_efdt_ratio` inmo_diffusion_nml.f90 + self.hdiff_efdt_ratio: float = hdiff_efdt_ratio + + #: Ratio of e-folding time to time step for w diffusion (NH only) + #: Called `hdiff_w_efdt_ratio` inmo_diffusion_nml.f90. + self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio + + #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below + #: Called `hdiff_smag_fac` inmo_diffusion_nml.f90 + self.smagorinski_scaling_factor: float = smagorinski_scaling_factor + + #: If True, apply truly horizontal temperature diffusion over steep slopes + #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90 + self.apply_zdiffusion_t: bool = zdiffu_t + + # from other namelists: + # from parent namelist mo_nonhydrostatic_nml + + #: Number of dynamics substeps per fast-physics step + #: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90 + self.ndyn_substeps: int = n_substeps + + #: If True, compute horizontal diffusion only at the large time step + #: Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90 + self.lhdiff_rcf: bool = hdiff_rcf + + # namelist mo_gridref_nml.f90 + + #: Denominator for temperature boundary diffusion + #: Called 'denom_diffu_t' in mo_gridref_nml.f90 + self.temperature_boundary_diffusion_denominator: float = ( + temperature_boundary_diffusion_denom + ) + + #: Denominator for velocity boundary diffusion + #: Called 'denom_diffu_v' in mo_gridref_nml.f90 + self.velocity_boundary_diffusion_denominator: float = velocity_boundary_diffusion_denom + + # parameters from namelist: mo_interpol_nml.f90 + + #: Parameter describing the lateral boundary nudging in limited area mode. + #: + #: Maximal value of the nudging coefficients used cell row bordering the boundary interpolation zone, + #: from there nudging coefficients decay exponentially with `nudge_efold_width` in units of cell rows. + #: Called `nudge_max_coeff` in mo_interpol_nml.f90 + self.nudge_max_coeff: float = max_nudging_coeff + + #: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients + #: Called `nudge_efold_width` in mo_interpol_nml.f90 + self.nudge_efold_width: float = nudging_decay_rate + + self._validate() + + def _validate(self): + """Apply consistency checks and validation on configuration parameters.""" + if self.diffusion_type != 5: + raise NotImplementedError( + "Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background " + "diffusion` is implemented" + ) + + if self.diffusion_type < 0: + self.apply_to_temperature = False + self.apply_to_horizontal_wind = False + self.apply_to_vertical_wind = False + else: + self.apply_to_temperature = True + self.apply_to_horizontal_wind = True + + if not self.apply_zdiffusion_t: + raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)") + + @functools.cached_property + def substep_as_float(self): + return float(self.ndyn_substeps) + + +@dataclass(frozen=True) +class DiffusionParams: + """Calculates derived quantities depending on the diffusion config.""" + + config: InitVar[DiffusionConfig] + K2: Final[float] = field(init=False) + K4: Final[float] = field(init=False) + K6: Final[float] = field(init=False) + K4W: Final[float] = field(init=False) + smagorinski_factor: Final[float] = field(init=False) + smagorinski_height: Final[float] = field(init=False) + scaled_nudge_max_coeff: Final[float] = field(init=False) + + def __post_init__(self, config): + object.__setattr__( + self, + "K2", + (1.0 / (config.hdiff_efdt_ratio * 8.0) if config.hdiff_efdt_ratio > 0.0 else 0.0), + ) + object.__setattr__(self, "K4", self.K2 / 8.0) + object.__setattr__(self, "K6", self.K2 / 64.0) + object.__setattr__( + self, + "K4W", + (1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0), + ) + + ( + smagorinski_factor, + smagorinski_height, + ) = self._determine_smagorinski_factor(config) + object.__setattr__(self, "smagorinski_factor", smagorinski_factor) + object.__setattr__(self, "smagorinski_height", smagorinski_height) + # see mo_interpol_nml.f90: + object.__setattr__( + self, + "scaled_nudge_max_coeff", + config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + ) + + def _determine_smagorinski_factor(self, config: DiffusionConfig): + """Enhanced Smagorinsky diffusion factor. + + Smagorinsky diffusion factor is defined as a profile in height + above sea level with 4 height sections. + + It is calculated/used only in the case of diffusion_type 3 or 5 + """ + match config.diffusion_type: + case 5: + ( + smagorinski_factor, + smagorinski_height, + ) = diffusion_type_5_smagorinski_factor(config) + case 4: + # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only + # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 + smagorinski_factor = ( + config.smagorinski_scaling_factor + if config.smagorinski_scaling_factor + else 0.15, + ) + smagorinski_height = None + case _: + raise NotImplementedError("Only implemented for diffusion type 4 and 5") + smagorinski_factor = None + smagorinski_height = None + pass + return smagorinski_factor, smagorinski_height + + +def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): + """ + Initialize Smagorinski factors used in diffusion type 5. + + The calculation and magic numbers are taken from mo_diffusion_nml.f90 + """ + magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) + magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) + magic_z2 = 1600.0 + 50000.0 + magic_sqrt + factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) + heights = (32500.0, magic_z2, 50000.0, 90000.0) + return factor, heights + + +class Diffusion: + """Class that configures diffusion and does one diffusion step.""" + + def __init__(self, exchange: ExchangeRuntime = SingleNode()): + self._exchange = exchange + self._initialized = False + self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) + self.thresh_tdiff: float = ( + -5.0 + ) #: threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + self.grid: Optional[IconGrid] = None + self.config: Optional[DiffusionConfig] = None + self.params: Optional[DiffusionParams] = None + self.vertical_params: Optional[VerticalModelParams] = None + self.interpolation_state: DiffusionInterpolationState = None + self.metric_state: DiffusionMetricState = None + self.diff_multfac_w: Optional[float] = None + self.diff_multfac_n2w: Field[[KDim], float] = None + self.smag_offset: Optional[float] = None + self.fac_bdydiff_v: Optional[float] = None + self.bdy_diff: Optional[float] = None + self.nudgezone_diff: Optional[float] = None + self.edge_params: Optional[EdgeParams] = None + self.cell_params: Optional[CellParams] = None + self._horizontal_start_index_w_diffusion: int32 = 0 + + def init( + self, + grid: IconGrid, + config: DiffusionConfig, + params: DiffusionParams, + vertical_params: VerticalModelParams, + metric_state: DiffusionMetricState, + interpolation_state: DiffusionInterpolationState, + edge_params: EdgeParams, + cell_params: CellParams, + ): + """ + Initialize Diffusion granule with configuration. + + calculates all local fields that are used in diffusion within the time loop. + + Args: + grid: + config: + params: + vertical_params: + metric_state: + interpolation_state: + edge_params: + cell_params: + """ + self.config: DiffusionConfig = config + self.params: DiffusionParams = params + self.grid = grid + self.vertical_params = vertical_params + self.metric_state: DiffusionMetricState = metric_state + self.interpolation_state: DiffusionInterpolationState = interpolation_state + self.edge_params = edge_params + self.cell_params = cell_params + + self._allocate_temporary_fields() + + def _get_start_index_for_w_diffusion() -> int32: + return self.grid.get_start_index( + CellDim, + ( + HorizontalMarkerIndex.nudging(CellDim) + if self.grid.limited_area() + else HorizontalMarkerIndex.interior(CellDim) + ), + ) + + self.nudgezone_diff: float = 0.04 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) + self.bdy_diff: float = 0.015 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) + self.fac_bdydiff_v: float = ( + math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator + if config.lhdiff_rcf + else 1.0 / config.velocity_boundary_diffusion_denominator + ) + + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float + self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float) + + init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( + params.K4, + config.substep_as_float, + *params.smagorinski_factor, + *params.smagorinski_height, + self.vertical_params.physical_heights, + self.diff_multfac_vn, + self.smag_limit, + self.enh_smag_fac, + offset_provider={"Koff": KDim}, + ) + + self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( + k_size=self.grid.n_lev(), + nshift=0, + physical_heights=np.asarray(self.vertical_params.physical_heights), + nrdmax=self.vertical_params.index_of_damping_layer, + ) + self._horizontal_start_index_w_diffusion = _get_start_index_for_w_diffusion() + self._initialized = True + + @property + def initialized(self): + return self._initialized + + def _allocate_temporary_fields(self): + def _allocate(*dims: Dimension): + return zero_field(self.grid, *dims) + + def _index_field(dim: Dimension, size=None): + size = size if size else self.grid.size[dim] + return np_as_located_field(dim)(np.arange(size, dtype=int32)) + + self.diff_multfac_vn = _allocate(KDim) + + self.smag_limit = _allocate(KDim) + self.enh_smag_fac = _allocate(KDim) + self.u_vert = _allocate(VertexDim, KDim) + self.v_vert = _allocate(VertexDim, KDim) + self.kh_smag_e = _allocate(EdgeDim, KDim) + self.kh_smag_ec = _allocate(EdgeDim, KDim) + self.z_nabla2_e = _allocate(EdgeDim, KDim) + self.z_temp = _allocate(CellDim, KDim) + self.diff_multfac_smag = _allocate(KDim) + self.z_nabla4_e2 = _allocate(EdgeDim, KDim) + # TODO(Magdalena): this is KHalfDim + self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) + self.horizontal_cell_index = _index_field(CellDim) + self.horizontal_edge_index = _index_field(EdgeDim) + self.w_tmp = np_as_located_field(CellDim, KDim)( + np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float) + ) + + def initial_run( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + ): + """ + Calculate initial diffusion step. + + In ICON at the start of the simulation diffusion is run with a parameter linit = True: + + 'For real-data runs, perform an extra diffusion call before the first time + step because no other filtering of the interpolated velocity field is done' + + This run uses special values for diff_multfac_vn, smag_limit and smag_offset + + """ + diff_multfac_vn = zero_field(self.grid, KDim) + smag_limit = zero_field(self.grid, KDim) + + setup_fields_for_initial_step.with_backend(backend)( + self.params.K4, + self.config.hdiff_efdt_ratio, + diff_multfac_vn, + smag_limit, + offset_provider={}, + ) + self._do_diffusion_step( + diagnostic_state, + prognostic_state, + dtime, + diff_multfac_vn, + smag_limit, + 0.0, + ) + self._sync_cell_fields(prognostic_state) + + def run( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + ): + """ + Do one diffusion step within regular time loop. + + runs a diffusion step for the parameter linit=False, within regular time loop. + """ + + self._do_diffusion_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + diff_multfac_vn=self.diff_multfac_vn, + smag_limit=self.smag_limit, + smag_offset=self.smag_offset, + ) + if not self.config.lhdiff_rcf: + self._sync_cell_fields(prognostic_state) + + def _sync_cell_fields(self, prognostic_state): + """ + Communicate theta_v, exner and w. + + communication only done in original code if the following condition applies: + IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + """ + log.debug("communication of prognostic cell fields: theta, w, exner - start") + handle_cell_comm = self._exchange.exchange( + CellDim, + prognostic_state.w, + prognostic_state.theta_v, + prognostic_state.exner_pressure, + ) + handle_cell_comm.wait() + log.debug("communication of prognostic cell fields: theta, w, exner - done") + + def _do_diffusion_step( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + dtime: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + ): + """ + Run a diffusion step. + + Args: + diagnostic_state: output argument, data class that contains diagnostic variables + prognostic_state: output argument, data class that contains prognostic variables + dtime: the time step, + diff_multfac_vn: + smag_limit: + smag_offset: + + """ + klevels = self.grid.n_lev() + cell_start_interior = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) + ) + cell_start_nudging = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.nudging(CellDim) + ) + cell_end_local = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) + cell_end_halo = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.halo(CellDim)) + + edge_start_nudging_plus_one = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 + ) + edge_start_nudging = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + ) + edge_start_lb_plus4 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + ) + edge_end_local = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.local(EdgeDim)) + edge_end_local_minus2 = self.grid.get_end_index( + EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 + ) + edge_end_halo = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.halo(EdgeDim)) + + vertex_start_lb_plus1 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 + ) + vertex_end_local = self.grid.get_end_index( + VertexDim, HorizontalMarkerIndex.local(VertexDim) + ) + + # dtime dependent: enh_smag_factor, + scale_k.with_backend(backend)( + self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} + ) + + log.debug("rbf interpolation 1: start") + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + p_e_in=prognostic_state.vn, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + ) + log.debug("rbf interpolation 1: end") + + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert + log.debug("communication rbf extrapolation of vn - start") + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h.wait() + log.debug("communication rbf extrapolation of vn - end") + + log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( + diff_multfac_smag=self.diff_multfac_smag, + tangent_orientation=self.edge_params.tangent_orientation, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_x=self.edge_params.primal_normal_vert[0], + primal_normal_vert_y=self.edge_params.primal_normal_vert[1], + dual_normal_vert_x=self.edge_params.dual_normal_vert[0], + dual_normal_vert_y=self.edge_params.dual_normal_vert[1], + vn=prognostic_state.vn, + smag_limit=smag_limit, + kh_smag_e=self.kh_smag_e, + kh_smag_ec=self.kh_smag_ec, + z_nabla2_e=self.z_nabla2_e, + smag_offset=smag_offset, + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local_minus2, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start") + calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( + kh_smag_ec=self.kh_smag_ec, + vn=prognostic_state.vn, + e_bln_c_s=self.interpolation_state.e_bln_c_s, + geofac_div=self.interpolation_state.geofac_div, + diff_multfac_smag=self.diff_multfac_smag, + wgtfac_c=self.metric_state.wgtfac_c, + div_ic=diagnostic_state.div_ic, + hdef_ic=diagnostic_state.hdef_ic, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=1, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), + "Koff": KDim, + }, + ) + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end") + + # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + + if self.config.type_vn_diffu > 1: + log.debug("communication rbf extrapolation of z_nable2_e - start") + h_z = self._exchange.exchange(EdgeDim, self.z_nabla2_e) + h_z.wait() + log.debug("communication rbf extrapolation of z_nable2_e - end") + + log.debug("2nd rbf interpolation: start") + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + p_e_in=self.z_nabla2_e, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + ) + log.debug("2nd rbf interpolation: end") + + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) + log.debug("communication rbf extrapolation of z_nable2_e - start") + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h.wait() + log.debug("communication rbf extrapolation of z_nable2_e - end") + + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") + apply_diffusion_to_vn.with_backend(backend)( + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], + primal_normal_vert_v2=self.edge_params.primal_normal_vert[1], + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + area_edge=self.edge_params.edge_areas, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + vn=prognostic_state.vn, + horz_idx=self.horizontal_edge_index, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v=self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + limited_area=self.grid.limited_area(), + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") + log.debug("communication of prognistic.vn : start") + handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) + + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + ) + # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify? + copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( + area=self.cell_params.area, + geofac_n2s=self.interpolation_state.geofac_n2s, + geofac_grg_x=self.interpolation_state.geofac_grg_x, + geofac_grg_y=self.interpolation_state.geofac_grg_y, + w_old=self.w_tmp, + w=prognostic_state.w, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=self.vertical_index, + horz_idx=self.horizontal_cell_index, + nrdmax=int32( + self.vertical_params.index_of_damping_layer + 1 + ), # +1 since Fortran includes boundaries + interior_idx=int32(cell_start_interior), + halo_idx=int32(cell_end_local), + horizontal_start=self._horizontal_start_index_w_diffusion, + horizontal_end=cell_end_halo, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity(), + }, + ) + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + ) + + log.debug( + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + ) + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(backend)( + theta_v=prognostic_state.theta_v, + theta_ref_mc=self.metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + horizontal_start=edge_start_nudging, + horizontal_end=edge_end_halo, + vertical_start=(klevels - 2), + vertical_end=klevels, + offset_provider={ + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + }, + ) + log.debug( + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + ) + log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") + calculate_nabla2_for_theta.with_backend(backend)( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths, + theta_v=prognostic_state.theta_v, + geofac_div=self.interpolation_state.geofac_div, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "E2C": self.grid.get_e2c_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), + }, + ) + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") + log.debug( + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + ) + truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)( + mask=self.metric_state.mask_hdiff, + zd_vertoffset=self.metric_state.zd_vertoffset, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + vcoef=self.metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2CEC": self.grid.get_c2cec_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "Koff": KDim, + }, + ) + + log.debug( + "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + ) + log.debug("running stencil 16 (update_theta_and_exner): start") + update_theta_and_exner.with_backend(backend)( + z_temp=self.z_temp, + area=self.cell_params.area, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={}, + ) + log.debug("running stencil 16 (update_theta_and_exner): end") + handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. + log.debug("communication of prognogistic.vn - end") diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py new file mode 100644 index 000000000..ea423047c --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -0,0 +1,115 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import functools +from dataclasses import dataclass + +import numpy as np +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.common.dimension import ( + C2E2CODim, + CECDim, + CEDim, + CellDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) + + +@dataclass(frozen=True) +class DiffusionDiagnosticState: + """Represents the diagnostic fields needed in diffusion.""" + + # fields for 3D elements in turbdiff + hdef_ic: Field[ + [CellDim, KDim], float + ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] + div_ic: Field[ + [CellDim, KDim], float + ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] + dwdx: Field[ + [CellDim, KDim], float + ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + + dwdy: Field[ + [CellDim, KDim], float + ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) + + +@dataclass(frozen=True) +class DiffusionMetricState: + """Represents the metric state fields needed in diffusion.""" + + theta_ref_mc: Field[[CellDim, KDim], float] + wgtfac_c: Field[ + [CellDim, KDim], float + ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) + mask_hdiff: Field[[CellDim, KDim], bool] + zd_vertoffset: Field[[CECDim, KDim], int32] + zd_diffcoef: Field[[CellDim, KDim], float] + zd_intcoef: Field[[CECDim, KDim], float] + + +@dataclass(frozen=True) +class DiffusionInterpolationState: + """Represents the ICON interpolation state needed in diffusion.""" + + e_bln_c_s: Field[[CEDim], float] # coefficent for bilinear interpolation from edge to cell () + rbf_coeff_1: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + + geofac_div: Field[[CEDim], float] # factor for divergence (nproma,cell_type,nblks_c) + + geofac_n2s: Field[ + [CellDim, C2E2CODim], float + ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) + geofac_grg_x: Field[[CellDim, C2E2CODim], float] + geofac_grg_y: Field[ + [CellDim, C2E2CODim], float + ] # factors for green gauss gradient (nproma,4,nblks_c,2) + nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges + + @functools.cached_property + def geofac_n2s_c(self) -> Field[[CellDim], float]: + return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) + + @functools.cached_property + def geofac_n2s_nbh(self) -> Field[[CECDim], float]: + geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:] + old_shape = geofac_nbh_ar.shape + return np_as_located_field(CECDim)( + geofac_nbh_ar.reshape( + old_shape[0] * old_shape[1], + ) + ) + + +@dataclass +class PrognosticState: + """Class that contains the prognostic state. + + Corresponds to ICON t_nh_prog + """ + + w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py new file mode 100644 index 000000000..183d146cb --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -0,0 +1,230 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from typing import Tuple + +import numpy as np +from gt4py.next.common import Dimension, Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, minimum +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim + + +# TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py +def zero_field(mesh, *dims: Dimension, dtype=float): + shapex = tuple(map(lambda x: mesh.size[x], dims)) + return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) + + +@field_operator +def _identity_c_k(field: Field[[CellDim, KDim], float]) -> Field[[CellDim, KDim], float]: + return field + + +@program +def copy_field(old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float]): + _identity_c_k(old_f, out=new_f) + + +@field_operator +def _identity_e_k(field: Field[[EdgeDim, KDim], float]) -> Field[[EdgeDim, KDim], float]: + return field + + +@field_operator +def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: + return field * factor + + +@program +def scale_k(field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float]): + _scale_k(field, factor, out=scaled_field) + + +@field_operator +def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: + return broadcast(0.0, (VertexDim, KDim)) + + +@program +def set_zero_v_k(field: Field[[VertexDim, KDim], float]): + _set_zero_v_k(out=field) + + +@field_operator +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: + return 0.125 - 4.0 * diff_multfac_vn + + +@field_operator +def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float) -> Field[[KDim], float]: + con = 1.0 / 128.0 + dyn = k4 * dyn_substeps / 3.0 + return broadcast(minimum(con, dyn), (KDim,)) + + +@field_operator +def _setup_initial_diff_multfac_vn(k4: float, hdiff_efdt_ratio: float) -> Field[[KDim], float]: + return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) + + +@field_operator +def _setup_fields_for_initial_step( + k4: float, hdiff_efdt_ratio: float +) -> Tuple[Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) + smag_limit = _setup_smag_limit(diff_multfac_vn) + return diff_multfac_vn, smag_limit + + +@program +def setup_fields_for_initial_step( + k4: float, + hdiff_efdt_ratio: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], +): + _setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit)) + + +@field_operator +def _en_smag_fac_for_zero_nshift( + vect_a: Field[[KDim], float], + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, +) -> Field[[KDim], float]: + dz21 = hdiff_smag_z2 - hdiff_smag_z + alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 + df32 = hdiff_smag_fac3 - hdiff_smag_fac2 + df42 = hdiff_smag_fac4 - hdiff_smag_fac2 + dz32 = hdiff_smag_z3 - hdiff_smag_z2 + dz42 = hdiff_smag_z4 - hdiff_smag_z2 + + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (vect_a + vect_a(Koff[1])) + + dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z)) + dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2)) + enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) + return enh_smag_fac + + +@field_operator +def _init_diffusion_local_fields_for_regular_timestemp( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], +) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) + smag_limit = _setup_smag_limit(diff_multfac_vn) + enh_smag_fac = _en_smag_fac_for_zero_nshift( + vect_a, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + ) + return ( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ) + + +@program +def init_diffusion_local_fields_for_regular_timestep( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + enh_smag_fac: Field[[KDim], float], +): + _init_diffusion_local_fields_for_regular_timestemp( + k4, + dyn_substeps, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + vect_a, + out=( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ), + ) + + +def init_nabla2_factor_in_upper_damping_zone( + k_size: int, nrdmax: int32, nshift: int, physical_heights: np.ndarray +) -> Field[[KDim], float]: + """ + Calculate diff_multfac_n2w. + + numpy version, since gt4py does not allow non-constant indexing into fields + + Args + k_size: number of vertical levels + nrdmax: index of the level where rayleigh dampint starts + nshift: + physcial_heights: vector of physical heights [m] of the height levels + """ + # TODO(Magdalena): fix with as_offset in gt4py + + buffer = np.zeros(k_size) + buffer[1 : nrdmax + 1] = ( + 1.0 + / 12.0 + * ( + ( + physical_heights[1 + nshift : nrdmax + 1 + nshift] + - physical_heights[nshift + nrdmax + 1] + ) + / (physical_heights[1] - physical_heights[nshift + nrdmax + 1]) + ) + ** 4 + ) + return np_as_located_field(KDim)(buffer) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py similarity index 86% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py index 61560cb35..bceeecf44 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py @@ -14,12 +14,18 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import _calculate_nabla2_of_theta -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( _truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.model.atmosphere.dycore.update_theta_and_exner import _update_theta_and_exner +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + _update_theta_and_exner, +) from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py similarity index 85% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 333c21116..36fded768 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -10,20 +10,20 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( _apply_nabla2_and_nabla4_global_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.model.atmosphere.dycore.calculate_nabla4 import _calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -91,7 +91,7 @@ def _apply_diffusion_to_vn( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], @@ -110,6 +110,10 @@ def apply_diffusion_to_vn( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, limited_area: bool, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_vn( u_vert, @@ -130,4 +134,8 @@ def apply_diffusion_to_vn( start_2nd_nudge_line_idx_e, limited_area, out=vn, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py similarity index 83% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 294eb2ed1..dd4f0a264 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,14 +14,16 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( _calculate_horizontal_gradients_for_turbulence, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import _calculate_nabla2_for_w +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( + _calculate_nabla2_for_w, +) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim @@ -91,6 +93,10 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( nrdmax: int32, interior_idx: int32, halo_idx: int32, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area, @@ -108,4 +114,8 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( interior_idx, halo_idx, out=(w, dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py similarity index 75% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py index f8d62fe06..9e181ee32 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -39,5 +39,20 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w) + _apply_nabla2_to_w( + area, + z_nabla2_c, + geofac_n2s, + w, + diff_multfac_w, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py similarity index 74% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py index 42128fd8a..657bdd427 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim @@ -35,5 +35,19 @@ def apply_nabla2_to_w_in_upper_damping_layer( diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _apply_nabla2_to_w_in_upper_damping_layer( + w, + diff_multfac_n2w, + cell_area, + z_nabla2_c, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py similarity index 71% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py index e051aba45..6b497c45a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py @@ -12,23 +12,23 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, ) -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator def _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: @@ -43,12 +43,16 @@ def _calculate_diagnostic_quantities_for_turbulence( def calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec, @@ -58,4 +62,8 @@ def calculate_diagnostic_quantities_for_turbulence( diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py similarity index 72% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 49eec5e7c..6a997cd21 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -12,12 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim @@ -43,7 +43,19 @@ def calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( - theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + theta_v, + theta_ref_mc, + thresh_tdiff, + kh_smag_e, + out=kh_smag_e, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py index 5a0a98d83..d174d72b0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -36,5 +36,18 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_horizontal_gradients_for_turbulence(w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy)) + _calculate_horizontal_gradients_for_turbulence( + w, + geofac_grg_x, + geofac_grg_y, + out=(dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py similarity index 67% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py index cca6712de..d72cecd1d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py @@ -12,10 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import _calculate_nabla2_of_theta +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( + _calculate_nabla2_for_z, +) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim @@ -38,5 +42,19 @@ def calculate_nabla2_for_theta( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_nabla2_for_theta(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _calculate_nabla2_for_theta( + kh_smag_e, + inv_dual_edge_length, + theta_v, + geofac_div, + out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py similarity index 74% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py index cdba6c34f..c40fab441 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -31,5 +31,17 @@ def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c) + _calculate_nabla2_for_w( + w, + geofac_n2s, + out=z_nabla2_c, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py similarity index 77% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index a90aaddea..cb528e6fa 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -15,19 +15,19 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator def _temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s, axis=C2EDim) / diff_multfac_smag - div = neighbor_sum(vn(C2E) * geofac_div, axis=C2EDim) + kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div @@ -35,8 +35,8 @@ def _temporary_fields_for_turbulence_diagnostics( def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], kh_c: Field[[CellDim, KDim], float], div: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 92% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index fd189e669..55eb44f2a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -65,6 +65,10 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( vcoef: Field[[CECDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, @@ -76,4 +80,8 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( theta_v, z_temp, out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py index bc53f6897..b74d32485 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim @@ -39,5 +39,20 @@ def update_theta_and_exner( theta_v: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], rd_o_cvd: float, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _update_theta_and_exner(z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner)) + _update_theta_and_exner( + z_temp, + area, + theta_v, + exner, + rd_o_cvd, + out=(theta_v, exner), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py new file mode 100644 index 000000000..7ba065d2c --- /dev/null +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -0,0 +1,24 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +""" +Initialize pytest. + +Workaround for pytest not discovering those configuration function, when they are added to the +diffusion_test/conftest.py folder +""" +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index b6b255adc..a29c25ea6 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py_common>=0.0.5" + "icon4py-common>=0.0.5" ] description = "ICON dynamical core." dynamic = ['version'] diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 7ee115147..9e9274b36 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -10,24 +10,23 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pytest -from gt4py.next.program_processors.runners.roundtrip import executor -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh - -BACKENDS = {"embedded": executor} -MESHES = {"simple_mesh": SimpleMesh()} - - -@pytest.fixture( - ids=MESHES.keys(), - params=MESHES.values(), +from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + backend, + damping_height, + data_provider, + datapath, + download_ser_data, + grid_savepoint, + icon_grid, + linit, + mesh, + step_date_exit, + step_date_init, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, ) -def mesh(request): - return request.param - - -@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) -def backend(request): - return request.param diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index 8c4a1b154..053013639 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -27,7 +27,11 @@ class TestMoSolveNonhydroStencil39(StencilTest): @staticmethod def reference( - mesh, e_bln_c_s: np.array, z_w_concorr_me: np.array, wgtfac_c: np.array, **kwargs + mesh, + e_bln_c_s: np.array, + z_w_concorr_me: np.array, + wgtfac_c: np.array, + **kwargs, ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index c8c67b999..6774dc946 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -27,7 +27,11 @@ class TestMoSolveNonhydroStencil40(StencilTest): @staticmethod def reference( - mesh, e_bln_c_s: np.array, z_w_concorr_me: np.array, wgtfacq_c: np.array, **kwargs + mesh, + e_bln_c_s: np.array, + z_w_concorr_me: np.array, + wgtfacq_c: np.array, + **kwargs, ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) diff --git a/model/common/.pre-commit-config.yaml b/model/common/.pre-commit-config.yaml index 3d45afc46..0a905ba43 100644 --- a/model/common/.pre-commit-config.yaml +++ b/model/common/.pre-commit-config.yaml @@ -63,6 +63,16 @@ repos: rev: v1.3.0 hooks: - id: yesqa + additional_dependencies: + - flake8==4.0.1 + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments - repo: https://github.com/psf/black rev: '22.3.0' diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 0a90a8cb9..22a32d8ac 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -22,7 +22,10 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1" + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "netcdf4>=1.6.0" ] description = "Shared code for the icon4py model." dynamic = ['version'] diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py new file mode 100644 index 000000000..1c64d9f04 --- /dev/null +++ b/model/common/src/icon4py/model/common/constants.py @@ -0,0 +1,31 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Final + + +#: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), +#: see https://glossary.ametsoc.org/wiki/Gas_constant. +GAS_CONSTANT_DRY_AIR: Final[float] = 287.04 + +#: Specific heat at constant pressure [J/K/kg] +CPD: Final[float] = 1004.64 + +#: Gas constant for water vapor [J/K/kg], rv in ICON. +GAS_CONSTANT_WATER_VAPOR: Final[float] = 461.51 + +#: Av. gravitational acceleration [m/s^2] +GRAVITATIONAL_ACCELERATION: Final[float] = 9.8066 + +#: Default physics to dynamics time step ratio +DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0 diff --git a/model/common/src/icon4py/model/common/decomposition/__init__.py b/model/common/src/icon4py/model/common/decomposition/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/common/src/icon4py/model/common/decomposition/decomposed.py b/model/common/src/icon4py/model/common/decomposition/decomposed.py new file mode 100644 index 000000000..23b278b65 --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/decomposed.py @@ -0,0 +1,251 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import logging +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +import ghex +import ghex.unstructured as unstructured +import numpy as np +import numpy.ma as ma +from gt4py.next import Dimension + +from icon4py.model.common.decomposition.parallel_setup import ProcessProperties +from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim +from icon4py.model.common.utils import builder + + +log = logging.getLogger(__name__) + + +class DomainDescriptorIdGenerator: + _counter = 0 + _roundtrips = 0 + + def __init__(self, parallel_props: ProcessProperties): + self._comm_size = parallel_props.comm_size + self._roundtrips = parallel_props.rank + self._base = self._roundtrips * self._comm_size + + def __call__(self): + next_id = self._base + self._counter + if self._counter + 1 >= self._comm_size: + self._roundtrips = self._roundtrips + self._comm_size + self._base = self._roundtrips * self._comm_size + self._counter = 0 + else: + self._counter = self._counter + 1 + return next_id + + +class DecompositionInfo: + class EntryType(int, Enum): + ALL = (0,) + OWNED = (1,) + HALO = 2 + + @builder + def with_dimension(self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray): + masked_global_index = ma.array(global_index, mask=owner_mask) + self._global_index[dim] = masked_global_index + + def __init__(self, klevels: int): + self._global_index = {} + self._klevels = klevels + + @property + def klevels(self): + return self._klevels + + def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return self._to_local_index(dim) + case DecompositionInfo.EntryType.HALO: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[~mask] + case DecompositionInfo.EntryType.OWNED: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[mask] + + def _to_local_index(self, dim): + data = ma.getdata(self._global_index[dim], subok=False) + assert data.ndim == 1 + return np.arange(data.shape[0]) + + def owner_mask(self, dim: Dimension) -> np.ndarray: + return self._global_index[dim].mask + + def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return ma.getdata(self._global_index[dim], subok=False) + case DecompositionInfo.EntryType.OWNED: + global_index = self._global_index[dim] + return ma.getdata(global_index[global_index.mask]) + case DecompositionInfo.EntryType.HALO: + global_index = self._global_index[dim] + return ma.getdata(global_index[~global_index.mask]) + case _: + raise NotImplementedError() + + +class ExchangeResult(Protocol): + def wait(self): + ... + + def is_ready(self) -> bool: + ... + + +class ExchangeRuntime(Protocol): + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + ... + + def get_size(self): + ... + + def my_rank(self): + ... + + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +def create_exchange(props: ProcessProperties, decomp_info: DecompositionInfo) -> ExchangeRuntime: + """ + Create an Exchange depending on the runtime size. + + Depending on the number of processor a SingleNode version is returned or a GHEX context created and a Multinode returned. + """ + if props.comm_size > 1: + return GHexMultiNode(props, decomp_info) + else: + return SingleNode() + + +@dataclass +class SingleNode: + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + return SingleNodeResult() + + def my_rank(self): + return 0 + + def get_size(self): + return 1 + + +class SingleNodeResult: + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +class GHexMultiNode: + def __init__(self, props: ProcessProperties, domain_decomposition: DecompositionInfo): + self._context = ghex.context(ghex.mpi_comm(props.comm), True) + self._domain_id_gen = DomainDescriptorIdGenerator(props) + self._decomposition_info = domain_decomposition + self._domain_descriptors = { + CellDim: self._create_domain_descriptor( + CellDim, + ), + VertexDim: self._create_domain_descriptor( + VertexDim, + ), + EdgeDim: self._create_domain_descriptor(EdgeDim), + } + log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") + + self._patterns = { + CellDim: self._create_pattern(CellDim), + VertexDim: self._create_pattern(VertexDim), + EdgeDim: self._create_pattern(EdgeDim), + } + log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") + self._comm = unstructured.make_co(self._context) + log.info("communication object initialized") + + def _domain_descriptor_info(self, descr): + return f" domain_descriptor=[id='{descr.domain_id()}', size='{descr.size()}', inner_size='{descr.inner_size()}' (halo size='{descr.size() - descr.inner_size()}')" + + def get_size(self): + return self._context.size() + + def my_rank(self): + return self._context.rank() + + def _create_domain_descriptor(self, dim: Dimension): + all_global = self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) + local_halo = self._decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) + # first arg is the domain ID which builds up an MPI Tag. + # if those ids are not different for all domain descriptors the system might deadlock + # if two parallel exchanges with the same domain id are done + domain_desc = unstructured.domain_descriptor( + self._domain_id_gen(), all_global.tolist(), local_halo.tolist() + ) + log.debug( + f"domain descriptor for dim='{dim.value}' with properties {self._domain_descriptor_info(domain_desc)} created" + ) + return domain_desc + + def _create_pattern(self, horizontal_dim: Dimension): + assert horizontal_dim.kind == DimensionKind.HORIZONTAL + + global_halo_idx = self._decomposition_info.global_index( + horizontal_dim, DecompositionInfo.EntryType.HALO + ) + halo_generator = unstructured.halo_generator_with_gids(global_halo_idx) + log.debug(f"halo generator for dim='{horizontal_dim.value}' created") + pattern = unstructured.make_pattern( + self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] + ) + log.debug( + f"pattern for dim='{horizontal_dim.value}' and {self._domain_descriptor_info(self._domain_descriptors[horizontal_dim])} created" + ) + return pattern + + def exchange(self, dim: Dimension, *fields: tuple): + assert dim in [CellDim, EdgeDim, VertexDim] + pattern = self._patterns[dim] + assert pattern is not None, f"pattern for {dim.value} not found" + domain_descriptor = self._domain_descriptors[dim] + assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" + applied_patterns = [ + pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields + ] + handle = self._comm.exchange(applied_patterns) + log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") + return MultiNodeResult(handle, applied_patterns) + + +@dataclass +class MultiNodeResult: + handle: ... + pattern_refs: ... + + def wait(self): + self.handle.wait() + del self.pattern_refs + + def is_ready(self) -> bool: + return self.handle.is_ready() diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py new file mode 100644 index 000000000..11d5eaeb6 --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -0,0 +1,98 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import functools +import logging +from dataclasses import dataclass +from typing import Optional, Union + +import mpi4py +from mpi4py.MPI import Comm + + +mpi4py.rc.initialize = False + +CommId = Union[int, Comm, None] +log = logging.getLogger(__name__) + + +def get_processor_properties(with_mpi=False, comm_id: CommId = None): + def _get_current_comm_or_comm_world(comm_id: CommId) -> Comm: + if isinstance(comm_id, int): + comm = Comm.f2py(comm_id) + elif isinstance(comm_id, Comm): + comm = comm_id + else: + comm = mpi4py.MPI.COMM_WORLD + return comm + + if with_mpi: + init_mpi() + current_comm = _get_current_comm_or_comm_world(comm_id) + return ProcessProperties.from_mpi_comm(current_comm) + else: + return ProcessProperties.from_single_node() + + +def init_mpi(): + from mpi4py import MPI + + if not MPI.Is_initialized(): + log.info("initializing MPI") + MPI.Init() + + +def finalize_mpi(): + from mpi4py import MPI + + if not MPI.Is_finalized(): + log.info("finalizing MPI") + MPI.Finalize() + + +@dataclass(frozen=True) +class ProcessProperties: + comm: Optional[mpi4py.MPI.Comm] = None + + @functools.cached_property + def rank(self): + return self.comm.Get_rank() if self.comm else 0 + + @functools.cached_property + def comm_name(self): + return self.comm.Get_name() if self.comm else "" + + @functools.cached_property + def comm_size(self): + return self.comm.Get_size() if self.comm else 1 + + @classmethod + def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): + return ProcessProperties(comm) + + @classmethod + def from_single_node(cls): + return ProcessProperties() + + +class ParallelLogger(logging.Filter): + def __init__(self, process_properties: ProcessProperties = None): + super().__init__() + self._rank_info = "" + if process_properties and process_properties.comm_size > 1: + self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " + + def filter( # noqa: A003 # overwriting logging.Filter.filter() + self, record: logging.LogRecord + ) -> bool: + record.rank = self._rank_info + return True diff --git a/model/common/src/icon4py/model/common/dimension.py b/model/common/src/icon4py/model/common/dimension.py index a0be9cc72..fd6bb6111 100644 --- a/model/common/src/icon4py/model/common/dimension.py +++ b/model/common/src/icon4py/model/common/dimension.py @@ -28,7 +28,9 @@ E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) V2CDim = Dimension("V2C", DimensionKind.LOCAL) +C2VDim = Dimension("C2V", DimensionKind.LOCAL) V2EDim = Dimension("V2E", DimensionKind.LOCAL) +V2E2VDim = Dimension("V2E2V", DimensionKind.LOCAL) E2C2VDim = Dimension("E2C2V", DimensionKind.LOCAL) C2E2CODim = Dimension("C2E2CO", DimensionKind.LOCAL) E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL) @@ -37,6 +39,7 @@ E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim)) C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim)) V2C = FieldOffset("V2C", source=CellDim, target=(VertexDim, V2CDim)) +C2V = FieldOffset("C2V", source=VertexDim, target=(CellDim, C2VDim)) V2E = FieldOffset("V2E", source=EdgeDim, target=(VertexDim, V2EDim)) E2V = FieldOffset("E2V", source=VertexDim, target=(EdgeDim, E2VDim)) C2CE = FieldOffset("C2CE", source=CEDim, target=(CellDim, C2EDim)) diff --git a/model/common/src/icon4py/model/common/grid/__init__.py b/model/common/src/icon4py/model/common/grid/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py new file mode 100644 index 000000000..79382c354 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -0,0 +1,405 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import dataclasses +import logging +from enum import Enum +from typing import Optional +from uuid import UUID + +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from netCDF4 import Dataset + +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + C2VDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + EdgeDim, + V2CDim, + V2E2VDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize + + +class GridFileName(str, Enum): + pass + + +@dataclasses.dataclass +class GridFileField: + name: GridFileName + shape: tuple[int, ...] + + +def _validate_shape(data: np.array, field_definition: GridFileField): + if data.shape != field_definition.shape: + raise IconGridError( + f"invalid grid file field {field_definition.name} does not have dimension {field_definition.shape}" + ) + + +class GridFile: + """Represent and ICON netcdf grid file.""" + + INVALID_INDEX = -1 + + class PropertyName(GridFileName): + GRID_ID = "uuidOfHGrid" + PARENT_GRID_ID = "uuidOfParHGrid" + + class OffsetName(GridFileName): + """Names for connectivities used in the grid file.""" + + # e2c2e/e2c2eO: diamond edges (including origin) not present in grid file-> calculate? + # e2c2v: diamond vertices: not present in grid file -> constructed from e2c and c2v + + #: name of C2E2C connectivity in grid file: dims(nv=3, cell) + C2E2C = "neighbor_cell_index" + + #: name of V2E2V connectivity in gridfile: dims(ne=6, vertex), + #: all vertices of a pentagon/hexagon, same as V2C2V + V2E2V = "vertices_of_vertex" # does not exist in simple_mesh.py + + #: name of V2E dimension in grid file: dims(ne=6, vertex) + V2E = "edges_of_vertex" + + #: name fo V2C connectivity in grid file: dims(ne=6, vertex) + V2C = "cells_of_vertex" + + #: name of E2V connectivity in grid file: dims(nc=2, edge) + E2V = "edge_vertices" + + #: name of C2V connectivity in grid file: dims(nv=3, cell) + C2V = "vertex_of_cell" # does not exist in simple_mesh.py + + #: name of E2C connectivity in grid file: dims(nc=2, edge) + E2C = "adjacent_cell_of_edge" + + #: name of C2E connectivity in grid file: dims(nv=3, cell) + C2E = "edge_of_cell" + + class DimensionName(GridFileName): + """Dimension values (sizes) used in grid file.""" + + #: number of vertices + VERTEX_NAME = "vertex" + + #: number of edges + EDGE_NAME = "edge" + + #: number of cells + CELL_NAME = "cell" + + #: number of edges in a diamond: 4 + DIAMOND_EDGE_SIZE = "no" + + #: number of edges/cells neibhboring one vertex: 6 (for regular, non pentagons) + NEIGHBORS_TO_VERTEX_SIZE = "ne" + + #: number of cells edges, vertices and cells neighboring a cell: 3 + NEIGHBORS_TO_CELL_SIZE = "nv" + + #: number of vertices/cells neighboring an edge: 2 + NEIGHBORS_TO_EDGE_SIZE = "nc" + + #: number of child domains (for nesting) + MAX_CHILD_DOMAINS = "max_chdom" + + #: Grid refinement: maximal number in grid-refinement (refin_ctl) array for each dimension + CELL_GRF = "cell_grf" + EDGE_GRF = "edge_grf" + VERTEX_GRF = "vert_grf" + + class GridRefinementName(GridFileName): + """Names of arrays in grid file defining the grid control, definition of boundaries layers, start and end indices of horizontal zones.""" + + #: refine control value of cell indices + CONTROL_CELLS = "refin_c_ctrl" + + #: refine control value of edge indices + CONTROL_EDGES = "refin_e_ctrl" + + #: refine control value of vertex indices + CONTROL_VERTICES = "refin_v_ctrl" + + #: start indices of horizontal grid zones for cell fields + START_INDEX_CELLS = "start_idx_c" + + #: start indices of horizontal grid zones for edge fields + START_INDEX_EDGES = "start_idx_e" + + #: start indices of horizontal grid zones for vertex fields + START_INDEX_VERTICES = "start_idx_v" + + #: end indices of horizontal grid zones for cell fields + END_INDEX_CELLS = "end_idx_c" + + #: end indices of horizontal grid zones for edge fields + END_INDEX_EDGES = "end_idx_e" + + #: end indices of horizontal grid zones for vertex fields + END_INDEX_VERTICES = "end_idx_v" + + def __init__(self, dataset: Dataset): + self._dataset = dataset + self._log = logging.getLogger(__name__) + + def dimension(self, name: GridFileName) -> int: + return self._dataset.dimensions[name].size + + def int_field(self, name: GridFileName, transpose=True, dtype=np.int32) -> np.ndarray: + try: + nc_variable = self._dataset.variables[name] + + self._log.debug(f"reading {name}: {nc_variable}") + data = nc_variable[:] + data = np.array(data, dtype=dtype) + return np.transpose(data) if transpose else data + except KeyError: + msg = f"{name} does not exist in dataset" + self._log.warning(msg) + raise IconGridError(msg) + + +class IconGridError(RuntimeError): + pass + + +class IndexTransformation: + def get_offset_for_index_field( + self, + array: np.ndarray, + ): + return np.zeros(array.shape, dtype=np.int32) + + +class ToGt4PyTransformation(IndexTransformation): + def get_offset_for_index_field(self, array: np.ndarray): + """ + Calculate the index offset needed for usage with python. + + Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for + INVALID values which are marked with -1 in the grid file and are kept such. + """ + return np.where(array == GridFile.INVALID_INDEX, 0, -1) + + +class GridManager: + """ + Read ICON grid file and set up IconGrid. + + Reads an ICON grid file and extracts connectivity arrays and start-, end-indices for horizontal + domain boundaries. Provides an IconGrid instance for further usage. + """ + + def __init__( + self, + transformation: IndexTransformation, + grid_file: str, + config: VerticalGridSize, + ): + self._log = logging.getLogger(__name__) + self._transformation = transformation + self._config = config + self._grid: Optional[IconGrid] = None + self._file_name = grid_file + + def __call__(self): + dataset = self._read_gridfile(self._file_name) + _, grid = self._constuct_grid(dataset) + self._grid = grid + + def _read_gridfile(self, fname: str) -> Dataset: + try: + dataset = Dataset(self._file_name, "r", format="NETCDF4") + self._log.debug(dataset) + return dataset + except FileNotFoundError: + self._log.error(f"gridfile {fname} not found, aborting") + exit(1) + + def _read_grid_refinement_information(self, dataset): + _CHILD_DOM = 0 + reader = GridFile(dataset) + + refin_ctrl = { + CellDim: reader.int_field(GridFile.GridRefinementName.CONTROL_CELLS), + EdgeDim: reader.int_field(GridFile.GridRefinementName.CONTROL_EDGES), + VertexDim: reader.int_field(GridFile.GridRefinementName.CONTROL_VERTICES), + } + refin_ctrl_max = { + CellDim: reader.dimension(GridFile.DimensionName.CELL_GRF), + EdgeDim: reader.dimension(GridFile.DimensionName.EDGE_GRF), + VertexDim: reader.dimension(GridFile.DimensionName.VERTEX_GRF), + } + start_indices = { + CellDim: self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_EDGES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_VERTICES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + } + end_indices = { + CellDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_CELLS, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_EDGES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_VERTICES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + } + + return start_indices, end_indices, refin_ctrl, refin_ctrl_max + + def get_grid(self): + return self._grid + + def _get_index(self, dim: Dimension, start_marker: int, index_dict): + if dim.kind != DimensionKind.HORIZONTAL: + msg = f"getting start index in horizontal domain with non - horizontal dimension {dim}" + self._log.warning(msg) + raise IconGridError(msg) + try: + return index_dict[dim][start_marker] + except KeyError: + msg = f"start, end indices for dimension {dim} not present" + self._log.error(msg) + raise IconGridError(msg) + + def _constuct_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: + grid_id = UUID(dataset.getncattr(GridFile.PropertyName.GRID_ID)) + return grid_id, self._from_grid_dataset(dataset) + + def get_size(self, dim: Dimension): + if dim == VertexDim: + return self._grid.config.num_vertices + elif dim == CellDim: + return self._grid.config.num_cells + elif dim == EdgeDim: + return self._grid.config.num_edges + else: + self._log.warning(f"cannot determine size of unknown dimension {dim}") + raise IconGridError(f"Unknown dimension {dim}") + + def _get_index_field( + self, + reader, + field: GridFileName, + transpose=True, + apply_offset=True, + dtype=np.int32, + ): + field = reader.int_field(field, transpose=transpose, dtype=dtype) + if apply_offset: + field = field + self._transformation.get_offset_for_index_field(field) + return field + + def _from_grid_dataset(self, dataset: Dataset) -> IconGrid: + reader = GridFile(dataset) + num_cells = reader.dimension(GridFile.DimensionName.CELL_NAME) + num_edges = reader.dimension(GridFile.DimensionName.EDGE_NAME) + num_vertices = reader.dimension(GridFile.DimensionName.VERTEX_NAME) + + grid_size = HorizontalGridSize( + num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells + ) + c2e = self._get_index_field(reader, GridFile.OffsetName.C2E) + + e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) + c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) + e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) + + e2c2v = self._construct_diamond_array(c2v, e2c) + + v2c = self._get_index_field(reader, GridFile.OffsetName.V2C) + v2e = self._get_index_field(reader, GridFile.OffsetName.V2E) + v2e2v = self._get_index_field(reader, GridFile.OffsetName.V2E2V) + c2e2c = self._get_index_field(reader, GridFile.OffsetName.C2E2C) + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + ( + start_indices, + end_indices, + refine_ctrl, + refine_ctrl_max, + ) = self._read_grid_refinement_information(dataset) + + config = GridConfig( + horizontal_config=grid_size, + vertical_config=self._config, + ) + icon_grid = ( + IconGrid() + .with_config(config) + .with_connectivities( + { + C2EDim: c2e, + E2CDim: e2c, + E2VDim: e2v, + V2EDim: v2e, + V2CDim: v2c, + C2VDim: c2v, + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + E2C2VDim: e2c2v, + V2E2VDim: v2e2v, + } + ) + .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim]) + .with_start_end_indices(EdgeDim, start_indices[EdgeDim], end_indices[EdgeDim]) + .with_start_end_indices(VertexDim, start_indices[VertexDim], end_indices[VertexDim]) + ) + + return icon_grid + + def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): + dummy_c2v = np.append( + c2v, + GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), + axis=0, + ) + expanded = dummy_c2v[e2c[:, :], :] + sh = expanded.shape + flattened = expanded.reshape(sh[0], sh[1] * sh[2]) + return np.apply_along_axis(np.unique, 1, flattened) diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py new file mode 100644 index 000000000..a68c95d1b --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -0,0 +1,269 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass +from typing import Final + +from gt4py.next.common import Dimension, Field + +from icon4py.model.common.dimension import CellDim, ECVDim, EdgeDim, VertexDim + + +class HorizontalMarkerIndex: + """ + Handles constants indexing into the start_index and end_index fields. + + ICON uses a double indexing scheme for field indices marking the start and end of special + grid zone: The constants defined here (from mo_impl_constants.f90 and mo_impl_constants_grf.f90) + are the indices that are used to index into the start_idx and end_idx arrays + provided by the grid file where for each dimension the start index of the horizontal + "zones" are defined: + f.ex. an inlined access of the field F: Field[[CellDim], double] at the starting point of the lateral boundary zone would be + + F[start_idx_c[_LATERAL_BOUNDARY_CELLS] + + ICON uses a custom index range from [ICON_INDEX_OFFSET... ] such that the index 0 marks the + internal entities for _all_ dimensions (Cell, Edge, Vertex) that is why we define these + additional INDEX_OFFSETs here in order to swap back to a 0 base python array. + """ + + NUM_GHOST_ROWS: Final[int] = 2 + # values from mo_impl_constants.f90 + _ICON_INDEX_OFFSET_CELLS: Final[int] = 8 + _GRF_BOUNDARY_WIDTH_CELL: Final[int] = 4 + _MIN_RL_CELL_INT: Final[int] = -4 + _MIN_RL_CELL: Final[int] = _MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS + _MAX_RL_CELL: Final[int] = 5 + + _ICON_INDEX_OFFSET_VERTEX: Final[int] = 7 + _MIN_RL_VERTEX_INT: Final[int] = _MIN_RL_CELL_INT + _MIN_RL_VERTEX: Final[int] = _MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1) + _MAX_RL_VERTEX: Final[int] = _MAX_RL_CELL + + _ICON_INDEX_OFFSET_EDGES: Final[int] = 13 + _GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 + _MIN_RL_EDGE_INT: Final[int] = 2 * _MIN_RL_CELL_INT + _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) + _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL + + _LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES + _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES + _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES + _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES + _END_EDGES: Final[int] = 0 + + _LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS + _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS + _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS + _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS + _END_CELLS: Final[int] = 0 + + _LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX + _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX + _NUDGING_VERTICES: Final[int] = 0 + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX + _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX + _END_VERTICES: Final[int] = 0 + + _lateral_boundary = { + CellDim: _LATERAL_BOUNDARY_CELLS, + EdgeDim: _LATERAL_BOUNDARY_EDGES, + VertexDim: _LATERAL_BOUNDARY_VERTICES, + } + _local = { + CellDim: _LOCAL_CELLS, + EdgeDim: _LOCAL_EDGES, + VertexDim: _LOCAL_VERTICES, + } + _halo = { + CellDim: _HALO_CELLS, + EdgeDim: _HALO_EDGES, + VertexDim: _HALO_VERTICES, + } + _interior = { + CellDim: _INTERIOR_CELLS, + EdgeDim: _INTERIOR_EDGES, + VertexDim: _INTERIOR_VERTICES, + } + _nudging = { + CellDim: _NUDGING_CELLS, + EdgeDim: _NUDGING_EDGES, + VertexDim: _NUDGING_VERTICES, + } + _end = { + CellDim: _END_CELLS, + EdgeDim: _END_EDGES, + VertexDim: _END_VERTICES, + } + + @classmethod + def lateral_boundary(cls, dim: Dimension) -> int: + """Indicate lateral boundary. + + These points correspond to the sorted points in ICON, the marker can be incremented in order + to access higher order boundary lines + """ + return cls._lateral_boundary[dim] + + @classmethod + def local(cls, dim: Dimension) -> int: + """Indicate points that are owned by the processing unit, i.e. no halo points.""" + return cls._local[dim] + + @classmethod + def halo(cls, dim: Dimension) -> int: + return cls._halo[dim] + + @classmethod + def nudging(cls, dim: Dimension) -> int: + """Indicate the nudging zone.""" + return cls._nudging[dim] + + @classmethod + def interior(cls, dim: Dimension) -> int: + """Indicate interior i.e. unordered prognostic cells in ICON.""" + return cls._interior[dim] + + @classmethod + def end(cls, dim: Dimension) -> int: + return cls._end[dim] + + +@dataclass(frozen=True) +class HorizontalGridSize: + num_vertices: int + num_edges: int + num_cells: int + + +# TODO(Magdalena): allow initialization with only partial values +# (becomes tedious for testing otherwise): hence this should +# that should not be a data class +class EdgeParams: + def __init__( + self, + tangent_orientation=None, + primal_edge_lengths=None, + inverse_primal_edge_lengths=None, + dual_edge_lengths=None, + inverse_dual_edge_lengths=None, + inverse_vertex_vertex_lengths=None, + primal_normal_vert_x=None, + primal_normal_vert_y=None, + dual_normal_vert_x=None, + dual_normal_vert_y=None, + edge_areas=None, + ): + + self.tangent_orientation: Field[[EdgeDim], float] = tangent_orientation + r""" + Orientation of vector product of the edge and the adjacent cell centers + v3 + / \ + / \ + / c1 \ + / | \ + v1---|--->v2 + \ | / + \ v / + \ c2 / + \ / + v4 + +1 or -1 depending on whether the vector product of + (v2-v1) x (c2-c1) points outside (+) or inside (-) the sphere + + defined in ICON in mo_model_domain.f90:t_grid_edges%tangent_orientation + """ + + self.primal_edge_lengths: Field[[EdgeDim], float] = primal_edge_lengths + """ + Length of the triangle edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length + """ + + self.inverse_primal_edge_lengths: Field[[EdgeDim], float] = inverse_primal_edge_lengths + """ + Inverse of the triangle edge length: 1.0/primal_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_primal_edge_length + """ + + self.dual_edge_lengths: Field[[EdgeDim], float] = dual_edge_lengths + """ + Length of the hexagon/pentagon edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length + """ + + self.inverse_dual_edge_lengths: Field[[EdgeDim], float] = inverse_dual_edge_lengths + """ + Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length + """ + + self.inverse_vertex_vertex_lengths: Field[[EdgeDim], float] = inverse_vertex_vertex_lengths + r""" + Inverse distance between outer vertices of adjacent cells. + + v1-------- + | /| + | / | + | e | + | / | + |/ | + --------v2 + + inverse_vertex_vertex_length(e) = 1.0/|v2-v1| + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length + """ + + self.primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + primal_normal_vert_x, + primal_normal_vert_y, + ) + """ + Normal of the triangle edge, projected onto the location of the vertices + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert + """ + + self.dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + dual_normal_vert_x, + dual_normal_vert_y, + ) + """ + Tangent to the triangle edge, projected onto the location of vertices. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_normal_vert + """ + + self.edge_areas: Field[[EdgeDim], float] = edge_areas + """ + Area of the quadrilateral (two triangles) adjacent to the edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%area_edge + """ + + +@dataclass(frozen=True) +class CellParams: + area: Field[[CellDim], float] + """ + Area of a cell. + + defined int ICON in mo_model_domain.f90:t_grid_cells%area + """ diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py new file mode 100644 index 000000000..f514ca128 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -0,0 +1,197 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass +from typing import Dict + +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import NeighborTableOffsetProvider +from typing_extensions import deprecated + +from icon4py.model.common.dimension import CECDim, CEDim, CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.horizontal import HorizontalGridSize +from icon4py.model.common.grid.vertical import VerticalGridSize +from icon4py.model.common.utils import builder + + +@dataclass(frozen=True) +class GridConfig: + horizontal_config: HorizontalGridSize + vertical_config: VerticalGridSize + limited_area: bool = True + n_shift_total: int = 0 + + @property + def num_k_levels(self): + return self.vertical_config.num_lev + + @property + def num_vertices(self): + return self.horizontal_config.num_vertices + + @property + def num_edges(self): + return self.horizontal_config.num_edges + + @property + def num_cells(self): + return self.horizontal_config.num_cells + + +class IconGrid: + def __init__(self): + self.config: GridConfig = None + + self.start_indices = {} + self.end_indices = {} + self.connectivities: Dict[str, np.ndarray] = {} + self.size: Dict[Dimension, int] = {} + + def _update_size(self, config: GridConfig): + self.size[VertexDim] = config.num_vertices + self.size[CellDim] = config.num_cells + self.size[EdgeDim] = config.num_edges + self.size[KDim] = config.num_k_levels + + @builder + def with_config(self, config: GridConfig): + self.config = config + self._update_size(config) + + @builder + def with_start_end_indices( + self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray + ): + self.start_indices[dim] = start_indices.astype(int32) + self.end_indices[dim] = end_indices.astype(int32) + + @builder + def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): + self.connectivities.update( + {d.value.lower(): k.astype(int) for d, k in connectivity.items()} + ) + self.size.update({d: t.shape[1] for d, t in connectivity.items()}) + + def limited_area(self): + # defined in mo_grid_nml.f90 + return self.config.limited_area + + def n_lev(self): + return self.config.num_k_levels if self.config else 0 + + def num_cells(self): + return self.config.num_cells if self.config else 0 + + def num_vertices(self): + return self.config.num_vertices if self.config else 0 + + def num_edges(self): + return self.config.num_edges + + @deprecated( + "use get_start_index and get_end_index instead, - should be removed after merge of solve_nonhydro" + ) + def get_indices_from_to( + self, dim: Dimension, start_marker: int, end_marker: int + ) -> tuple[int32, int32]: + """ + Use to specify domains of a field for field_operator. + + For a given dimension, returns the start and end index if a + horizontal region in a field given by the markers. + + field operators will then run from start of the region given by the + start_marker to the end of the region given by the end_marker + """ + if dim.kind != DimensionKind.HORIZONTAL: + raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) + return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] + + def get_start_index(self, dim: Dimension, marker: int) -> int32: + """ + Use to specify lower end of domains of a field for field_operators. + + For a given dimension, returns the start index of the + horizontal region in a field given by the marker. + """ + return self.start_indices[dim][marker] + + def get_end_index(self, dim: Dimension, marker: int) -> int32: + """ + Use to specify upper end of domains of a field for field_operators. + + For a given dimension, returns the end index of the + horizontal region in a field given by the marker. + """ + return self.end_indices[dim][marker] + + def get_c2e_connectivity(self): + table = self.connectivities["c2e"] + return NeighborTableOffsetProvider(table, CellDim, EdgeDim, table.shape[1]) + + def get_e2c_connectivity(self): + table = self.connectivities["e2c"] + return NeighborTableOffsetProvider(table, EdgeDim, CellDim, table.shape[1]) + + def get_e2v_connectivity(self): + table = self.connectivities["e2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) + + def get_c2e2c_connectivity(self): + table = self.connectivities["c2e2c"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + def get_c2e2co_connectivity(self): + table = self.connectivities["c2e2co"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + def get_e2c2v_connectivity(self): + table = self.connectivities["e2c2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) + + def get_v2e_connectivity(self): + table = self.connectivities["v2e"] + return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + + def get_v2c_connectivity(self): + table = self.connectivities["v2c"] + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) + + def get_c2v_connectivity(self): + table = self.connectivities["c2v"] + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) + + def get_e2ecv_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["e2c2v"].shape, EdgeDim, ECVDim + ) + + def _neighbortable_offset_provider_for_1d_sparse_fields( + self, + old_shape: tuple[int, int], + origin_axis: Dimension, + neighbor_axis: Dimension, + ): + table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(table, origin_axis, neighbor_axis, table.shape[1]) + + def get_c2cec_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e2c"].shape, CellDim, CECDim + ) + + def get_c2ce_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e"].shape, CellDim, CEDim + ) diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py new file mode 100644 index 000000000..061e6a786 --- /dev/null +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -0,0 +1,56 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import Field, dataclass, field +from typing import Final + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.common.dimension import KDim + + +@dataclass(frozen=True) +class VerticalGridSize: + num_lev: int + + +@dataclass(frozen=True) +class VerticalModelParams: + """ + Contains vertical physical parameters defined on the grid. + + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml + """ + + vct_a: Field[[KDim], float] + rayleigh_damping_height: Final[float] + index_of_damping_layer: Final[int32] = field(init=False) + + def __post_init__(self): + object.__setattr__( + self, + "index_of_damping_layer", + self._determine_damping_height_index( + np.asarray(self.vct_a), self.rayleigh_damping_height + ), + ) + + @classmethod + def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): + return int32(np.argmax(np.where(vct_a >= damping_height))) + + @property + def physical_heights(self) -> Field[[KDim], float]: + return self.vct_a diff --git a/model/common/src/icon4py/model/common/interpolation/__init__.py b/model/common/src/icon4py/model/common/interpolation/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py new file mode 100644 index 000000000..cd7fa4ab7 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + + +def compute_c_lin_e( + edge_cell_length: np.array, + inv_dual_edge_length: np.array, + owner_mask: np.array, + second_boundary_layer_start_index: np.int32, +) -> np.array: + """ + Compute E2C average inverse distance. + + Args: + edge_cell_length: numpy array, representing a Field[[EdgeDim, E2CDim], float] + inv_dual_edge_length: inverse dual edge length, numpy array representing a Field[[EdgeDim], float] + owner_mask: numpy array, representing a Field[[EdgeDim], bool]boolean field, True for all edges owned by this compute node + second_boundary_layer_start_index: start index of the 2nd boundary line: c_lin_e is not calculated for the first boundary layer + + Returns: c_lin_e: numpy array representing Field[[EdgeDim, E2CDim], float] + + """ + c_lin_e_ = edge_cell_length[:, 1] * inv_dual_edge_length + c_lin_e = np.transpose(np.vstack((c_lin_e_, (1.0 - c_lin_e_)))) + c_lin_e[0:second_boundary_layer_start_index, :] = 0.0 + mask = np.transpose(np.tile(owner_mask, (2, 1))) + return np.where(mask, c_lin_e, 0.0) diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 76% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py index 0a9bf2d68..bf5bb4c3c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -36,5 +36,18 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) + _mo_intp_rbf_rbf_vec_interpol_vertex( + p_e_in, + ptr_coeff_1, + ptr_coeff_2, + out=(p_u_out, p_v_out), + domain={ + VertexDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/test_utils/data_handling.py b/model/common/src/icon4py/model/common/test_utils/data_handling.py new file mode 100644 index 000000000..7a8a49c34 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/data_handling.py @@ -0,0 +1,30 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import tarfile +from pathlib import Path + +import wget + + +def download_and_extract(uri: str, local_path: Path, data_file: str) -> None: + local_path.mkdir(parents=True, exist_ok=True) + if not any(local_path.iterdir()): + print(f"directory {local_path} is empty: downloading data from {uri} and extracting") + wget.download(uri, out=data_file) + # extract downloaded file + if not tarfile.is_tarfile(data_file): + raise NotImplementedError(f"{data_file} needs to be a valid tar file") + with tarfile.open(data_file, mode="r:*") as tf: + tf.extractall(path=local_path) + Path(data_file).unlink(missing_ok=True) diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py new file mode 100644 index 000000000..c6a52f0c9 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -0,0 +1,183 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from pathlib import Path + +import pytest +from gt4py.next.program_processors.runners.roundtrip import executor + +from ..decomposition.parallel_setup import get_processor_properties +from .data_handling import download_and_extract +from .serialbox_utils import IconSerialDataProvider +from .simple_mesh import SimpleMesh + + +test_utils = Path(__file__).parent +model = test_utils.parent.parent +common = model.parent.parent.parent.parent +base_path = common.parent.joinpath("testdata") + +data_uris = { + 1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download", + 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", + 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", +} + +ser_data_basepath = base_path.joinpath("ser_icondata") + + +@pytest.fixture(params=[False], scope="session") +def processor_props(request): + with_mpi = request.param + return get_processor_properties(with_mpi=with_mpi) + + +@pytest.fixture(scope="session") +def ranked_data_path(processor_props): + return ser_data_basepath.absolute().joinpath(f"mpitask{processor_props.comm_size}") + + +@pytest.fixture(scope="session") +def datapath(ranked_data_path): + return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") + + +@pytest.fixture(scope="session") +def download_ser_data(request, processor_props, ranked_data_path, pytestconfig): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + if not pytestconfig.getoption("datatest"): + pytest.skip("not running datatest marked tests") + + try: + uri = data_uris[processor_props.comm_size] + + data_file = ranked_data_path.joinpath( + f"mch_ch_r04b09_dsl_mpitask{processor_props.comm_size}.tar.gz" + ).name + if processor_props.rank == 0: + download_and_extract(uri, ranked_data_path, data_file) + if processor_props.comm: + processor_props.comm.barrier() + except KeyError: + raise AssertionError( + f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" + ) + + +@pytest.fixture(scope="session") +def data_provider(download_ser_data, datapath, processor_props) -> IconSerialDataProvider: + return IconSerialDataProvider( + fname_prefix="icon_pydycore", + path=str(datapath), + mpi_rank=processor_props.rank, + do_print=True, + ) + + +@pytest.fixture +def grid_savepoint(data_provider): + return data_provider.from_savepoint_grid() + + +@pytest.fixture +def icon_grid(grid_savepoint): + """ + Load the icon grid from an ICON savepoint. + + Uses the special grid_savepoint that contains data from p_patch + """ + return grid_savepoint.construct_icon_grid() + + +@pytest.fixture +def decomposition_info(data_provider): + return data_provider.from_savepoint_grid().construct_decomposition_info() + + +@pytest.fixture +def damping_height(): + return 12500 + + +@pytest.fixture +def ndyn_substeps(): + """ + Return number of dynamical substeps. + + Serialized data uses a reduced number (2 instead of the default 5) in order to reduce the amount + of data generated. + """ + return 2 + + +@pytest.fixture +def linit(): + """ + Set the 'linit' flag for the ICON diffusion data savepoint. + + Defaults to False + """ + return False + + +@pytest.fixture +def step_date_init(): + """ + Set the step date for the loaded ICON time stamp at start of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def step_date_exit(): + """ + Set the step date for the loaded ICON time stamp at the end of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def interpolation_savepoint(data_provider): # F811 + """Load data from ICON interplation state savepoint.""" + return data_provider.from_interpolation_savepoint() + + +@pytest.fixture +def metrics_savepoint(data_provider): # F811 + """Load data from ICON mestric state savepoint.""" + return data_provider.from_metrics_savepoint() + + +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py new file mode 100644 index 000000000..6dbe56db6 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -0,0 +1,22 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest + +from icon4py.model.common.decomposition.parallel_setup import ProcessProperties + + +def check_comm_size(props: ProcessProperties, sizes=(1, 2, 4)): + if props.comm_size not in sizes: + pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py new file mode 100644 index 000000000..de4425c8a --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/pytest_config.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + + +def pytest_configure(config): + config.addinivalue_line("markers", "datatest: this test uses binary data") + + +def pytest_addoption(parser): + """Add --datatest commandline option for pytest. + + Makes sure the option is set only once even when running tests of several model packages in one session. + """ + try: + parser.addoption( + "--datatest", + action="store_true", + help="running tests that use serialized data, can be slow since data might be downloaded from online storage", + default=False, + ) + except ValueError: + pass + + +def pytest_runtest_setup(item): + for _ in item.iter_markers(name="datatest"): + if not item.config.getoption("--datatest"): + pytest.skip("need '--datatest' option to run") diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py new file mode 100644 index 000000000..f4550ccb5 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -0,0 +1,577 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import logging + +import numpy as np +import serialbox as ser +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4py.model.atmosphere.diffusion.diffusion import VectorTuple +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.common import dimension +from icon4py.model.common.decomposition.decomposed import DecompositionInfo +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CECDim, + CEDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field + + +class IconSavepoint: + def __init__(self, sp: ser.Savepoint, ser: ser.Serializer, size: dict): + self.savepoint = sp + self.serializer = ser + self.sizes = size + self.log = logging.getLogger((__name__)) + + def log_meta_info(self): + self.log.info(self.savepoint.metainfo) + + def _get_field(self, name, *dimensions, dtype=float): + buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) + buffer_size = ( + self.sizes[d] if d.kind is DimensionKind.HORIZONTAL else s + for s, d in zip(buffer.shape, dimensions) + ) + buffer = buffer[tuple(map(slice, buffer_size))] + + self.log.debug(f"{name} {buffer.shape}") + return np_as_located_field(*dimensions)(buffer) + + def get_metadata(self, *names): + metadata = self.savepoint.metainfo.to_dict() + return {n: metadata[n] for n in names if n in metadata} + + def _read_int32_shift1(self, name: str): + """ + Read a start indices field. + + use for start indices: the shift accounts for the zero based python + values are converted to int32 + """ + return self._read_int32(name, offset=1) + + def _read_int32(self, name: str, offset=0): + """ + Read an end indices field. + + use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based + this accounts for being exclusive python exclusive bounds: [from:to) + field values are convert to int32 + """ + return self._read(name, offset, dtype=int32) + + def _read_bool(self, name: str): + return self._read(name, offset=0, dtype=bool) + + def _read(self, name: str, offset=0, dtype=int): + return (self.serializer.read(name, self.savepoint) - offset).astype(dtype) + + +class IconGridSavePoint(IconSavepoint): + def vct_a(self): + return self._get_field("vct_a", KDim) + + def tangent_orientation(self): + return self._get_field("tangent_orientation", EdgeDim) + + def inverse_primal_edge_lengths(self): + return self._get_field("inv_primal_edge_length", EdgeDim) + + def inv_vert_vert_length(self): + return self._get_field("inv_vert_vert_length", EdgeDim) + + def primal_normal_vert_x(self): + return self._get_field("primal_normal_vert_x", VertexDim, E2C2VDim) + + def primal_normal_vert_y(self): + return self._get_field("primal_normal_vert_y", VertexDim, E2C2VDim) + + def dual_normal_vert_y(self): + return self._get_field("dual_normal_vert_y", VertexDim, E2C2VDim) + + def dual_normal_vert_x(self): + return self._get_field("dual_normal_vert_x", VertexDim, E2C2VDim) + + def cell_areas(self): + return self._get_field("cell_areas", CellDim) + + def edge_areas(self): + return self._get_field("edge_areas", EdgeDim) + + def inv_dual_edge_length(self): + return self._get_field("inv_dual_edge_length", EdgeDim) + + def edge_cell_length(self): + return self._get_field("edge_cell_length", EdgeDim, E2CDim) + + def cells_start_index(self): + return self._read_int32_shift1("c_start_index") + + def cells_end_index(self): + return self._read_int32("c_end_index") + + def vertex_start_index(self): + return self._read_int32_shift1("v_start_index") + + def vertex_end_index(self): + return self._read_int32("v_end_index") + + def edge_start_index(self): + return self._read_int32_shift1("e_start_index") + + def edge_end_index(self): + # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being + # one off accounts for being exclusive [from:to) + return self.serializer.read("e_end_index", self.savepoint) + + def c_owner_mask(self): + return self._get_field("c_owner_mask", CellDim, dtype=bool) + + def e_owner_mask(self): + return self._get_field("e_owner_mask", EdgeDim, dtype=bool) + + def print_connectivity_info(self, name: str, ar: np.ndarray): + self.log.debug(f" connectivity {name} {ar.shape}") + + def c2e(self): + return self._get_connectivity_array("c2e", CellDim) + + def _get_connectivity_array(self, name: str, target_dim: Dimension): + connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :] + self.log.debug(f" connectivity {name} : {connectivity.shape}") + return connectivity + + def c2e2c(self): + return self._get_connectivity_array("c2e2c", CellDim) + + def e2c(self): + return self._get_connectivity_array("e2c", EdgeDim) + + def e2v(self): + # array "e2v" is actually e2c2v + v_ = self._get_connectivity_array("e2v", EdgeDim)[:, 0:2] + self.log.debug(f"real e2v {v_.shape}") + return v_ + + def e2c2v(self): + # array "e2v" is actually e2c2v, that is hexagon or pentagon + return self._get_connectivity_array("e2v", EdgeDim) + + def v2e(self): + return self._get_connectivity_array("v2e", VertexDim) + + def v2c(self): + return self._get_connectivity_array("v2c", VertexDim) + + def c2v(self): + return self._get_connectivity_array("c2v", CellDim) + + def nrdmax(self): + return self._read_int32_shift1("nrdmax") + + def refin_ctrl(self, dim: Dimension): + field_name = "refin_ctl" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def num(self, dim: Dimension): + return self.sizes[dim] + + def _read_field_for_dim(self, field_name, read_func, dim: Dimension): + match (dim): + case dimension.CellDim: + return read_func(f"c_{field_name}") + case dimension.EdgeDim: + return read_func(f"e_{field_name}") + case dimension.VertexDim: + return read_func(f"v_{field_name}") + case _: + raise NotImplementedError( + f"only {dimension.CellDim, dimension.EdgeDim, dimension.VertexDim} are handled" + ) + + def owner_mask(self, dim: Dimension): + field_name = "owner_mask" + mask = self._read_field_for_dim(field_name, self._read_bool, dim) + return np.squeeze(mask) + + def global_index(self, dim: Dimension): + field_name = "glb_index" + return self._read_field_for_dim(field_name, self._read_int32_shift1, dim) + + def decomp_domain(self, dim): + field_name = "decomp_domain" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def construct_decomposition_info(self): + return ( + DecompositionInfo(klevels=self.num(KDim)) + .with_dimension(*self._get_decomp_fields(CellDim)) + .with_dimension(*self._get_decomp_fields(EdgeDim)) + .with_dimension(*self._get_decomp_fields(VertexDim)) + ) + + def _get_decomp_fields(self, dim: Dimension): + global_index = self.global_index(dim) + mask = self.owner_mask(dim)[0 : self.num(dim)] + return dim, global_index, mask + + def construct_icon_grid(self) -> IconGrid: + + cell_starts = self.cells_start_index() + cell_ends = self.cells_end_index() + vertex_starts = self.vertex_start_index() + vertex_ends = self.vertex_end_index() + edge_starts = self.edge_start_index() + edge_ends = self.edge_end_index() + config = GridConfig( + horizontal_config=HorizontalGridSize( + num_vertices=self.num(VertexDim), + num_cells=self.num(CellDim), + num_edges=self.num(EdgeDim), + ), + vertical_config=VerticalGridSize(num_lev=self.num(KDim)), + ) + c2e2c = self.c2e2c() + c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c)) + grid = ( + IconGrid() + .with_config(config) + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) + .with_start_end_indices(EdgeDim, edge_starts, edge_ends) + .with_start_end_indices(CellDim, cell_starts, cell_ends) + .with_connectivities( + { + C2EDim: self.c2e(), + E2CDim: self.e2c(), + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + } + ) + .with_connectivities({E2VDim: self.e2v(), V2EDim: self.v2e(), E2C2VDim: self.e2c2v()}) + ) + return grid + + def construct_edge_geometry(self) -> EdgeParams: + primal_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.primal_normal_vert_y(), ECVDim), + ) + dual_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.dual_normal_vert_y(), ECVDim), + ) + return EdgeParams( + tangent_orientation=self.tangent_orientation(), + inverse_primal_edge_lengths=self.inverse_primal_edge_lengths(), + inverse_dual_edge_lengths=self.inv_dual_edge_length(), + inverse_vertex_vertex_lengths=self.inv_vert_vert_length(), + primal_normal_vert_x=primal_normal_vert[0], + primal_normal_vert_y=primal_normal_vert[1], + dual_normal_vert_x=dual_normal_vert[0], + dual_normal_vert_y=dual_normal_vert[1], + edge_areas=self.edge_areas(), + ) + + def construct_cell_geometry(self) -> CellParams: + return CellParams(area=self.cell_areas()) + + +class InterpolationSavepoint(IconSavepoint): + def geofac_grg(self): + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + num_cells = self.sizes[CellDim] + return np_as_located_field(CellDim, C2E2CODim)(grg[:num_cells, :, 0]), np_as_located_field( + CellDim, C2E2CODim + )(grg[:num_cells, :, 1]) + + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + + def e_bln_c_s(self): + return self._get_field("e_bln_c_s", CellDim, C2EDim) + + def geofac_div(self): + return self._get_field("geofac_div", CellDim, C2EDim) + + def geofac_n2s(self): + return self._get_field("geofac_n2s", CellDim, C2E2CODim) + + def rbf_vec_coeff_v1(self): + return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) + + def rbf_vec_coeff_v2(self): + return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + + def nudgecoeff_e(self): + return self._get_field("nudgecoeff_e", EdgeDim) + + def construct_interpolation_state_for_diffusion( + self, + ) -> DiffusionInterpolationState: + grg = self.geofac_grg() + return DiffusionInterpolationState( + e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), + rbf_coeff_1=self.rbf_vec_coeff_v1(), + rbf_coeff_2=self.rbf_vec_coeff_v2(), + geofac_div=as_1D_sparse_field(self.geofac_div(), CEDim), + geofac_n2s=self.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=self.nudgecoeff_e(), + ) + + def c_lin_e(self): + return self._get_field("c_lin_e", EdgeDim, E2CDim) + + +class MetricSavepoint(IconSavepoint): + def construct_metric_state_for_diffusion(self) -> DiffusionMetricState: + return DiffusionMetricState( + mask_hdiff=self.mask_diff(), + theta_ref_mc=self.theta_ref_mc(), + wgtfac_c=self.wgtfac_c(), + zd_intcoef=self.zd_intcoef(), + zd_vertoffset=self.zd_vertoffset(), + zd_diffcoef=self.zd_diffcoef(), + ) + + def zd_diffcoef(self): + return self._get_field("zd_diffcoef", CellDim, KDim) + + def zd_intcoef(self): + return self._read_and_reorder_sparse_field("vcoef", CellDim) + + def _read_and_reorder_sparse_field(self, name: str, horizontal_dim: Dimension, sparse_size=3): + ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[ + : self.sizes[horizontal_dim], :, : + ] + if ser_input.shape[1] != sparse_size: + ser_input = np.moveaxis((ser_input), 1, -1) + + return self._linearize_first_2dims( + ser_input, sparse_size=sparse_size, target_dims=(CECDim, KDim) + ) + + def _linearize_first_2dims( + self, data: np.ndarray, sparse_size: int, target_dims: tuple[Dimension, ...] + ): + old_shape = data.shape + assert old_shape[1] == sparse_size + return np_as_located_field(*target_dims)( + data.reshape(old_shape[0] * old_shape[1], old_shape[2]) + ) + + def zd_vertoffset(self): + return self._read_and_reorder_sparse_field("zd_vertoffset", CellDim) + + def zd_vertidx(self): + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + + def zd_indlist(self): + return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) + + def theta_ref_mc(self): + return self._get_field("theta_ref_mc", CellDim, KDim) + + def wgtfac_c(self): + return self._get_field("wgtfac_c", CellDim, KDim) + + def wgtfac_e(self): + return self._get_field("wgtfac_e", EdgeDim, KDim) + + def mask_diff(self): + return self._get_field("mask_hdiff", CellDim, KDim, dtype=bool) + + +class IconDiffusionInitSavepoint(IconSavepoint): + def hdef_ic(self): + return self._get_field("hdef_ic", CellDim, KDim) + + def div_ic(self): + return self._get_field("div_ic", CellDim, KDim) + + def dwdx(self): + return self._get_field("dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("dwdy", CellDim, KDim) + + def vn(self): + return self._get_field("vn", EdgeDim, KDim) + + def theta_v(self): + return self._get_field("theta_v", CellDim, KDim) + + def w(self): + return self._get_field("w", CellDim, KDim) + + def exner(self): + return self._get_field("exner", CellDim, KDim) + + def diff_multfac_smag(self): + return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) + + def smag_limit(self): + return np.squeeze(self.serializer.read("smag_limit", self.savepoint)) + + def diff_multfac_n2w(self): + return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) + + def nudgezone_diff(self) -> int: + return self.serializer.read("nudgezone_diff", self.savepoint)[0] + + def bdy_diff(self) -> int: + return self.serializer.read("bdy_diff", self.savepoint)[0] + + def fac_bdydiff_v(self) -> int: + return self.serializer.read("fac_bdydiff_v", self.savepoint)[0] + + def smag_offset(self): + return self.serializer.read("smag_offset", self.savepoint)[0] + + def diff_multfac_w(self): + return self.serializer.read("diff_multfac_w", self.savepoint)[0] + + def diff_multfac_vn(self): + return self.serializer.read("diff_multfac_vn", self.savepoint) + + def construct_prognostics(self) -> PrognosticState: + return PrognosticState( + w=self.w(), + vn=self.vn(), + exner_pressure=self.exner(), + theta_v=self.theta_v(), + ) + + def construct_diagnostics_for_diffusion(self) -> DiffusionDiagnosticState: + return DiffusionDiagnosticState( + hdef_ic=self.hdef_ic(), + div_ic=self.div_ic(), + dwdx=self.dwdx(), + dwdy=self.dwdy(), + ) + + +class IconDiffusionExitSavepoint(IconSavepoint): + def vn(self): + return self._get_field("x_vn", EdgeDim, KDim) + + def theta_v(self): + return self._get_field("x_theta_v", CellDim, KDim) + + def w(self): + return self._get_field("x_w", CellDim, KDim) + + def dwdx(self): + return self._get_field("x_dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("x_dwdy", CellDim, KDim) + + def exner(self): + return self._get_field("x_exner", CellDim, KDim) + + def z_temp(self): + return self._get_field("x_z_temp", CellDim, KDim) + + def div_ic(self): + return self._get_field("x_div_ic", CellDim, KDim) + + def hdef_ic(self): + return self._get_field("x_hdef_ic", CellDim, KDim) + + +class IconSerialDataProvider: + def __init__(self, fname_prefix, path=".", do_print=False, mpi_rank=0): + self.rank = mpi_rank + self.serializer: ser.Serializer = None + self.file_path: str = path + self.fname = f"{fname_prefix}_rank{str(self.rank)}" + self.log = logging.getLogger(__name__) + self._init_serializer(do_print) + self.grid_size = self._grid_size() + + def _init_serializer(self, do_print: bool): + if not self.fname: + self.log.warning(" WARNING: no filename! closing serializer") + self.serializer = ser.Serializer(ser.OpenModeKind.Read, self.file_path, self.fname) + if do_print: + self.print_info() + + def print_info(self): + self.log.info(f"SAVEPOINTS: {self.serializer.savepoint_list()}") + self.log.info(f"FIELDNAMES: {self.serializer.fieldnames()}") + + def _grid_size(self): + sp = self._get_icon_grid_savepoint() + grid_sizes = { + CellDim: self.serializer.read("num_cells", savepoint=sp).astype(int32)[0], + EdgeDim: self.serializer.read("num_edges", savepoint=sp).astype(int32)[0], + VertexDim: self.serializer.read("num_vert", savepoint=sp).astype(int32)[0], + KDim: sp.metainfo.to_dict()["nlev"], + } + return grid_sizes + + def from_savepoint_grid(self) -> IconGridSavePoint: + savepoint = self._get_icon_grid_savepoint() + return IconGridSavePoint(savepoint, self.serializer, size=self.grid_size) + + def _get_icon_grid_savepoint(self): + savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() + return savepoint + + def from_savepoint_diffusion_init( + self, + linit: bool, + date: str, + ) -> IconDiffusionInitSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-init"].linit[linit].date[date].as_savepoint() + ) + return IconDiffusionInitSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_interpolation_savepoint(self) -> InterpolationSavepoint: + savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() + return InterpolationSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_metrics_savepoint(self) -> MetricSavepoint: + savepoint = self.serializer.savepoint["metric_state"].as_savepoint() + return MetricSavepoint(savepoint, self.serializer, size=self.grid_size) + + def from_savepoint_diffusion_exit(self, linit: bool, date: str) -> IconDiffusionExitSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-exit"].linit[linit].date[date].as_savepoint() + ) + return IconDiffusionExitSavepoint(savepoint, self.serializer, size=self.grid_size) diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index 74211edd1..4715be0fb 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -59,6 +59,29 @@ @dataclass class SimpleMeshData: + c2v_table = np.asarray( + [ + [0, 1, 4], + [1, 2, 5], + [2, 0, 3], + [0, 3, 4], + [1, 4, 5], + [2, 5, 3], + [3, 4, 7], + [4, 5, 8], + [5, 3, 6], + [3, 6, 7], + [4, 7, 8], + [5, 8, 6], + [6, 7, 1], + [7, 8, 2], + [8, 6, 0], + [6, 0, 1], + [7, 1, 2], + [8, 2, 0], + ] + ) + e2c2v_table = np.asarray( [ [0, 1, 4, 6], # 0 @@ -354,6 +377,7 @@ class SimpleMesh: def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.diamond_arr = SimpleMeshData.diamond_table + self.c2v = SimpleMeshData.c2v_table self.e2c = SimpleMeshData.e2c_table self.e2v = SimpleMeshData.e2v_table self.c2e = SimpleMeshData.c2e_table @@ -373,6 +397,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_e2c2e = self.e2c2e.shape[1] self.n_e2c2v = self.e2c2v.shape[1] self.n_v2c = self.v2c.shape[1] + self.n_c2v = self.c2v.shape[1] self.n_v2e = self.v2e.shape[1] self.n_cells = self.c2e.shape[0] self.n_edges = 27 @@ -397,6 +422,9 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): ECVDim: self.n_edges * self.n_e2c2v, } + def get_c2v_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2v, VertexDim, CellDim, self.n_c2v) + def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e, CellDim, EdgeDim, self.n_c2e) @@ -427,6 +455,16 @@ def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) + def get_e2ecv_offset_provider(self): + old_shape = self.e2c2v.shape + e2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(e2ecv_table, EdgeDim, ECVDim, e2ecv_table.shape[1]) + + def get_c2ce_offset_provider(self): + old_shape = self.c2e.shape + c2ce_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider(c2ce_table, CellDim, CEDim, c2ce_table.shape[1]) + def get_offset_provider(self): return { "C2E": self.get_c2e_offset_provider(), @@ -439,8 +477,8 @@ def get_offset_provider(self): "E2C": self.get_e2c_offset_provider(), "E2V": self.get_e2v_offset_provider(), "E2C2V": self.get_e2c2v_offset_provider(), + "C2CE": self.get_c2ce_offset_provider(), "Koff": KDim, - "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, self.n_c2e), "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, self.n_e2c2v), "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, self.n_e2c), } diff --git a/model/common/src/icon4py/model/common/utils.py b/model/common/src/icon4py/model/common/utils.py new file mode 100644 index 000000000..403bf7a68 --- /dev/null +++ b/model/common/src/icon4py/model/common/utils.py @@ -0,0 +1,22 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +def builder(func): + """Use as decorator on builder functions.""" + + def wrapper(self, *args, **kwargs): + func(self, *args, **kwargs) + return self + + return wrapper diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py new file mode 100644 index 000000000..63571d684 --- /dev/null +++ b/model/common/tests/conftest.py @@ -0,0 +1,64 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import pytest + +from icon4py.model.common.test_utils.data_handling import download_and_extract +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 + backend, + base_path, + damping_height, + data_provider, + datapath, + decomposition_info, + download_ser_data, + grid_savepoint, + icon_grid, + interpolation_savepoint, + mesh, + processor_props, + ranked_data_path, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) + + +grids_path = base_path.joinpath("grids") +r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name +r02b04_global_grid_path = grids_path.joinpath("r02b04_global") +r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name + + +mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" + + +@pytest.fixture() +def r04b09_dsl_gridfile(get_grid_files): + return r04b09_dsl_grid_path.joinpath("grid.nc") + + +@pytest.fixture(scope="session") +def get_grid_files(pytestconfig): + """ + Get the grid files used for testing. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + if not pytestconfig.getoption("datatest"): + pytest.skip("not running datatest marked tests") + download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) + download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py new file mode 100644 index 000000000..b363461f0 --- /dev/null +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -0,0 +1,196 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.common.decomposition.decomposed import ( + DecompositionInfo, + DomainDescriptorIdGenerator, + GHexMultiNode, + SingleNode, + create_exchange, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size + + +""" +running tests with mpi: + +mpirun -np 2 python -m pytest -v --with-mpi tests/mpi_tests/test_parallel_setup.py + +mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ + + +""" + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_props(processor_props): + assert processor_props.comm + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + ("dim, owned, total"), + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), +) +def test_decomposition_info_masked( + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 +): + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) + my_total = total[my_rank] + my_owned = owned[my_rank] + assert all_indices.shape[0] == my_total + + owned_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.OWNED) + assert owned_indices.shape[0] == my_owned + + halo_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == my_total - my_owned + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + + +# @pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +def _assert_index_partitioning(all_indices, halo_indices, owned_indices): + owned_list = owned_indices.tolist() + halos_list = halo_indices.tolist() + all_list = all_indices.tolist() + assert set(owned_list) & set(halos_list) == set() + assert set(owned_list) & set(all_list) == set(owned_list) + assert set(halos_list) & set(all_list) == set(halos_list) + assert set(halos_list) | set(owned_list) == set(all_list) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + ("dim, owned, total"), + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), +) +@pytest.mark.mpi(min_size=2) +def test_decomposition_info_local_index( + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 +): + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) + my_total = total[my_rank] + my_owned = owned[my_rank] + + assert all_indices.shape[0] == my_total + assert np.array_equal(all_indices, np.arange(0, my_total)) + halo_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == my_total - my_owned + assert halo_indices.shape[0] < all_indices.shape[0] + assert np.alltrue(halo_indices <= np.max(all_indices)) + + owned_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.OWNED) + assert owned_indices.shape[0] == my_owned + assert owned_indices.shape[0] <= all_indices.shape[0] + assert np.alltrue(owned_indices <= np.max(all_indices)) + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("num", [1, 2, 3]) +def test_domain_descriptor_id_are_globally_unique(num, processor_props): + props = processor_props + size = props.comm_size + id_gen = DomainDescriptorIdGenerator(parallel_props=props) + id1 = id_gen() + assert id1 == props.comm_size * props.rank + assert id1 < props.comm_size * (props.rank + 1) + ids = [] + ids.append(id1) + for _ in range(1, num * size): + next_id = id_gen() + assert next_id > id1 + ids.append(next_id) + all_ids = props.comm.gather(ids, root=0) + if props.rank == 0: + all_ids = np.asarray(all_ids).flatten() + assert len(all_ids) == size * size * num + assert len(all_ids) == len(set(all_ids)) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_decomposition_info_matches_gridsize( + caplog, + download_ser_data, + decomposition_info, + icon_grid, + processor_props, +): # F811 + + check_comm_size(processor_props) + assert ( + decomposition_info.global_index( + dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL + ).shape[0] + == icon_grid.num_cells() + ) + assert ( + decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[0] + == icon_grid.num_vertices() + ) + assert ( + decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[0] + == icon_grid.num_edges() + ) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_create_multi_node_runtime_with_mpi( + download_ser_data, decomposition_info, processor_props +): # F811 + props = processor_props + exchange = create_exchange(props, decomposition_info) + if props.comm_size > 1: + assert isinstance(exchange, GHexMultiNode) + else: + assert isinstance(exchange, SingleNode) + + +@pytest.mark.parametrize("processor_props", [False], indirect=True) +def test_create_single_node_runtime_without_mpi(processor_props, decomposition_info): + props = processor_props + exchange = create_exchange(props, decomposition_info) + + assert isinstance(exchange, SingleNode) diff --git a/model/common/tests/mpi_tests/test_parallel_setup.py b/model/common/tests/mpi_tests/test_parallel_setup.py new file mode 100644 index 000000000..bcf714005 --- /dev/null +++ b/model/common/tests/mpi_tests/test_parallel_setup.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest +from mpi4py import MPI + +from icon4py.model.common.decomposition.parallel_setup import get_processor_properties, init_mpi + + +@pytest.mark.mpi +def test_parallel_properties_from_comm_world(): + props = get_processor_properties(with_mpi=True) + assert props.rank < props.comm_size + assert props.comm_name == "MPI_COMM_WORLD" + + +@pytest.mark.mpi(min_size=2) +def test_parallel_properties_from_mpi_comm(): + init_mpi() + world = MPI.COMM_WORLD + group = world.Get_group() + pair = group.Incl([0, 1]) + comm = world.Create(pair) + if comm != MPI.COMM_NULL: + comm.Set_name("my_comm") + props = get_processor_properties(with_mpi=True, comm_id=comm) + assert props.rank < props.comm_size + assert props.comm_size == 2 + assert props.comm_name == "my_comm" + + +def test_single_node_properties(): + props = get_processor_properties() + assert props.comm_size == 1 + assert props.rank == 0 + assert props.comm_name == "" diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py new file mode 100644 index 000000000..6f5f6c47a --- /dev/null +++ b/model/common/tests/test_grid_manager.py @@ -0,0 +1,521 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import logging +from uuid import uuid4 + +import numpy as np +import pytest +from netCDF4 import Dataset + +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.grid_manager import ( + GridFile, + GridFileName, + GridManager, + IndexTransformation, + ToGt4PyTransformation, +) +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.grid.icon_grid import VerticalGridSize +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +SIMPLE_MESH_NC = "simple_mesh_grid.nc" + + +@pytest.fixture +def simple_mesh_gridfile(tmp_path): + path = tmp_path.joinpath(SIMPLE_MESH_NC).absolute() + mesh = SimpleMesh() + dataset = Dataset(path, "w", format="NETCDF4") + dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) + dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=mesh.n_vertices) + + dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=mesh.n_edges) + dataset.createDimension(GridFile.DimensionName.CELL_NAME, size=mesh.n_cells) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, size=mesh.n_e2v) + dataset.createDimension(GridFile.DimensionName.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) + dataset.createDimension(GridFile.DimensionName.MAX_CHILD_DOMAINS, size=1) + # add dummy values for the grf dimensions + dataset.createDimension(GridFile.DimensionName.CELL_GRF, size=14) + dataset.createDimension(GridFile.DimensionName.EDGE_GRF, size=24) + dataset.createDimension(GridFile.DimensionName.VERTEX_GRF, size=13) + _add_to_dataset( + dataset, + np.zeros(mesh.n_edges), + GridFile.GridRefinementName.CONTROL_EDGES, + (GridFile.DimensionName.EDGE_NAME,), + ) + + _add_to_dataset( + dataset, + np.zeros(mesh.n_cells), + GridFile.GridRefinementName.CONTROL_CELLS, + (GridFile.DimensionName.CELL_NAME,), + ) + _add_to_dataset( + dataset, + np.zeros(mesh.n_vertices), + GridFile.GridRefinementName.CONTROL_VERTICES, + (GridFile.DimensionName.VERTEX_NAME,), + ) + + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, size=mesh.n_c2e) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, size=mesh.n_v2c) + + _add_to_dataset( + dataset, + mesh.c2e, + GridFile.OffsetName.C2E, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.e2c, + GridFile.OffsetName.E2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), + ) + _add_to_dataset( + dataset, + mesh.e2v, + GridFile.OffsetName.E2V, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.v2c, + GridFile.OffsetName.V2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), + ) + + _add_to_dataset( + dataset, + mesh.c2v, + GridFile.OffsetName.C2V, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + _add_to_dataset( + dataset, + np.zeros((mesh.n_vertices, 4), dtype=np.int32), + GridFile.OffsetName.V2E2V, + (GridFile.DimensionName.DIAMOND_EDGE_SIZE, GridFile.DimensionName.VERTEX_NAME), + ) + _add_to_dataset( + dataset, + mesh.v2e, + GridFile.OffsetName.V2E, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), + ) + _add_to_dataset( + dataset, + mesh.c2e2c, + GridFile.OffsetName.C2E2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, + ), + ) + + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) + dataset.close() + yield path + path.unlink() + + +def _add_to_dataset( + dataset: Dataset, + data: np.ndarray, + var_name: str, + dims: tuple[GridFileName, GridFileName], +): + var = dataset.createVariable(var_name, np.int32, dims) + var[:] = np.transpose(data)[:] + + +def test_gridparser_dimension(simple_mesh_gridfile): + data = Dataset(simple_mesh_gridfile, "r") + grid_parser = GridFile(data) + mesh = SimpleMesh() + assert grid_parser.dimension(GridFile.DimensionName.CELL_NAME) == mesh.n_cells + assert grid_parser.dimension(GridFile.DimensionName.VERTEX_NAME) == mesh.n_vertices + assert grid_parser.dimension(GridFile.DimensionName.EDGE_NAME) == mesh.n_edges + + +@pytest.mark.datatest +def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): + data = Dataset(r04b09_dsl_gridfile, "r") + grid_file = GridFile(data) + + assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num(CellDim) + assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num(EdgeDim) + assert grid_file.dimension(GridFile.DimensionName.VERTEX_NAME) == grid_savepoint.num(VertexDim) + + +def test_grid_parser_index_fields(simple_mesh_gridfile, caplog): + caplog.set_level(logging.DEBUG) + data = Dataset(simple_mesh_gridfile, "r") + mesh = SimpleMesh() + grid_parser = GridFile(data) + + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.C2E), mesh.c2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.E2C), mesh.e2c) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2E), mesh.v2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2C), mesh.v2c) + + +# TODO @magdalena add test cases for hexagon vertices v2e2v +# v2e2v: grid,??? + +# v2e: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + seralized_v2e = grid_savepoint.v2e()[0 : grid.num_vertices(), :] + # there are vertices at the boundary of a local domain or at a pentagon point that have less than + # 6 neighbors hence there are "Missing values" in the grid file + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(seralized_v2e) + assert has_invalid_index(grid.get_v2e_connectivity().table) + reset_invalid_index(seralized_v2e) + assert np.allclose(grid.get_v2e_connectivity().table, seralized_v2e) + + +# v2c: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_v2c = grid_savepoint.v2c()[0 : grid.num_vertices(), :] + # there are vertices that have less than 6 neighboring cells: either pentagon points or + # vertices at the boundary of the domain for a limited area mode + # hence in the grid file there are "missing values" + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(serialized_v2c) + assert has_invalid_index(grid.get_v2c_connectivity().table) + reset_invalid_index(serialized_v2c) + + assert np.allclose(grid.get_v2c_connectivity().table, serialized_v2c) + + +def reset_invalid_index(index_array: np.ndarray): + """ + Revert changes from mo_model_domimp_patches. + + Helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. + see: + # ! Checks for the pentagon case and moves dummy cells to end. + # ! The dummy entry is either set to 0 or duplicated from the last one + # SUBROUTINE move_dummies_to_end(array, array_size, max_connectivity, duplicate) + + After reading the grid file ICON moves all invalid indices (neighbors not existing in the + grid file) to the end of the neighbor list and replaces them with the "last valid neighbor index" + it is up to the user then to ensure that any coefficients in neighbor some multiplied with + these values are zero in order to "remove" them again from the sum. + + For testing we resubstitute those to the GridFile.INVALID_INDEX + Args: + index_array: the array where values the invalid values have to be reset + + Returns: an array where the spurious "last valid index" are replaced by GridFile.INVALID_INDEX + + """ + for i in range(0, index_array.shape[0]): + uq, index = np.unique(index_array[i, :], return_index=True) + index_array[i, max(index) + 1 :] = GridFile.INVALID_INDEX + + +# e2v: exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + + serialized_e2v = grid_savepoint.e2v()[0 : grid.num_edges(), :] + # all vertices in the system have to neighboring edges, there no edges that point nowhere + # hence this connectivity has no "missing values" in the grid file + assert not has_invalid_index(serialized_e2v) + assert not has_invalid_index(grid.get_e2v_connectivity().table) + assert np.allclose(grid.get_e2v_connectivity().table, serialized_e2v) + + +def has_invalid_index(ar: np.ndarray): + return np.any(np.where(ar == GridFile.INVALID_INDEX)) + + +# e2c : exists in serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_e2c = grid_savepoint.e2c()[0 : grid.num_edges(), :] + # there are edges at the boundary that have only one + # neighboring cell, there are "missing values" in the grid file + # and here they do not get substituted in the ICON preprocessing + assert has_invalid_index(serialized_e2c) + assert has_invalid_index(grid.get_e2c_connectivity().table) + assert np.allclose(grid.get_e2c_connectivity().table, serialized_e2c) + + +# c2e: serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + + serialized_c2e = grid_savepoint.c2e()[0 : grid.num_cells(), :] + # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the + # first place + # hence there are no "missing values" in the grid file + assert not has_invalid_index(serialized_c2e) + assert not has_invalid_index(grid.get_c2e_connectivity().table) + assert np.allclose(grid.get_c2e_connectivity().table, serialized_c2e) + + +# e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh +@pytest.mark.datatest +@pytest.mark.skip("does not directly exist in the grid file, needs to be constructed") +# TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell +def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile) + serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :] + assert has_invalid_index(serialized_e2c2e) + grid = gm.get_grid() + assert has_invalid_index(grid.get_e2c2e_connectivity().table) + assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) + + +# c2e2c: exists in serial, simple_mesh, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert np.allclose( + grid.get_c2e2c_connectivity().table, + grid_savepoint.c2e2c()[0 : grid.num_cells(), :], + ) + + +@pytest.mark.xfail +@pytest.mark.datatest +def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + # the "far" (adjacent to edge normal ) is not there. why? + # despite that: ordering is different + assert np.allclose( + grid.get_e2c2v_connectivity().table, + grid_savepoint.e2c2v()[0 : grid.num_edges(), :], + ) + + +@pytest.mark.datatest +def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + c2v = grid.get_c2v_connectivity().table + assert np.allclose(c2v, grid_savepoint.c2v()[0 : grid.num_cells(), :]) + + +def init_grid_manager(fname): + gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(65)) + gm() + return gm + + +@pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) +def test_grid_manager_getsize(simple_mesh_gridfile, dim, size, caplog): + caplog.set_level(logging.DEBUG) + gm = GridManager(IndexTransformation(), simple_mesh_gridfile, VerticalGridSize(num_lev=80)) + gm() + assert size == gm.get_size(dim) + + +def test_grid_manager_diamond_offset(simple_mesh_gridfile): + mesh = SimpleMesh() + gm = GridManager( + IndexTransformation(), + simple_mesh_gridfile, + VerticalGridSize(num_lev=mesh.k_level), + ) + gm() + grid = gm.get_grid() + assert np.allclose( + np.sort(grid.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) + ) + + +def test_gridmanager_given_file_not_found_then_abort(): + fname = "./unknown_grid.nc" + with pytest.raises(SystemExit) as error: + gm = GridManager(IndexTransformation(), fname, VerticalGridSize(num_lev=80)) + gm() + assert error.type == SystemExit + assert error.value == 1 + + +@pytest.mark.parametrize("size", [100, 1500, 20000]) +def test_gt4py_transform_offset_by_1_where_valid(size): + trafo = ToGt4PyTransformation() + input_field = np.random.randint(-1, size, (size,)) + offset = trafo.get_offset_for_index_field(input_field) + expected = np.where(input_field >= 0, -1, 0) + assert np.allclose(expected, offset) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), + ], +) +def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_start_index(dim, marker) == index + assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), + ], +) +def test_get_end_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_end_index(dim, marker) == index + assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker) diff --git a/model/common/tests/test_icon_grid.py b/model/common/tests/test_icon_grid.py new file mode 100644 index 000000000..83f20ee41 --- /dev/null +++ b/model/common/tests/test_icon_grid.py @@ -0,0 +1,400 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import pytest + +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex + + +@pytest.mark.datatest +# TODO(Magdalena) HorizontalMarkerIndex.local(dim) does not yield equivalent results form grid file +# and serialized data, why?. Serialized data has those strange -1 values +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), + ], +) +def test_horizontal_end_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_end_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "dim, marker, index", + [ + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), -1), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), -1), # ???? + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), -1), # ??? + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), + ], +) +def test_horizontal_start_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_start_index(dim, marker) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(CellDim), + HorizontalMarkerIndex.lateral_boundary(CellDim), + (0, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + (850, 1688), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + (1688, 2511), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + (2511, 3316), + ), + ( + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.interior(CellDim), + (4104, 20896), + ), + ( + HorizontalMarkerIndex.interior(CellDim) + 1, + HorizontalMarkerIndex.interior(CellDim) + 1, + (0, 850), + ), + ( + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.nudging(CellDim), + ( + 3316, + 4104, + ), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.local(CellDim), + (-1, 20896), + ), + ], +) +def test_horizontal_cell_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + CellDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + (428, 1278), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + (1278, 1700), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + (1700, 2538), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + (2538, 2954), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + (2954, 3777), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + (3777, 4184), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + (4184, 4989), + ), + ( + HorizontalMarkerIndex.interior(EdgeDim), + HorizontalMarkerIndex.interior(EdgeDim), + (6176, 31558), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ( + 4989, + 5387, + ), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + (5387, 6176), + ), + ( + HorizontalMarkerIndex.end(EdgeDim), + HorizontalMarkerIndex.end(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.local(EdgeDim), + HorizontalMarkerIndex.local(EdgeDim), + (-1, 31558), + ), + ], +) +def test_horizontal_edge_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + EdgeDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + (428, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + (850, 1266), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + (1266, 1673), + ), + ( + HorizontalMarkerIndex.interior(VertexDim), + HorizontalMarkerIndex.interior(VertexDim), + (2071, 10663), + ), + ( + HorizontalMarkerIndex.interior(VertexDim) + 1, + HorizontalMarkerIndex.interior(VertexDim) + 1, + (0, 428), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.halo(VertexDim), + HorizontalMarkerIndex.halo(VertexDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.local(VertexDim), + HorizontalMarkerIndex.local(VertexDim), + (-1, 10663), + ), + ], +) +def test_horizontal_vertex_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + VertexDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +def test_cross_check_marker_equivalences(icon_grid): + """Check actual equivalences of calculated markers.""" + # TODO(Magdalena): This should go away once we refactor these markers in a good way, such that no calculation need to be done with them anymore. + + assert icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.local(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) - 1, + ) == icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ) + assert icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.nudging(CellDim) - 1, + ) == icon_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + ) + assert icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 1, + ) == icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) + + assert icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + ) == icon_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ) + + +@pytest.mark.datatest +def test_grid_size(grid_savepoint): + assert 10663 == grid_savepoint.num(VertexDim) + assert 20896 == grid_savepoint.num(CellDim) + assert 31558 == grid_savepoint.num(EdgeDim) diff --git a/model/common/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py new file mode 100644 index 000000000..e4c8eb5b9 --- /dev/null +++ b/model/common/tests/test_interpolation_fields.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.common.dimension import EdgeDim +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.interpolation.interpolation_fields import compute_c_lin_e + + +@pytest.mark.datatest +def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() + edge_cell_length = grid_savepoint.edge_cell_length() + owner_mask = grid_savepoint.e_owner_mask() + c_lin_e_ref = interpolation_savepoint.c_lin_e() + lateral_boundary = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + c_lin_e = compute_c_lin_e( + np.asarray(edge_cell_length), + np.asarray(inv_dual_edge_length), + np.asarray(owner_mask), + lateral_boundary, + ) + + assert np.allclose(c_lin_e, c_lin_e_ref) diff --git a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 85% rename from model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 809e48b20..189beae08 100644 --- a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,11 +13,12 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( +from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -51,4 +52,8 @@ def input_data(self, mesh): ptr_coeff_2=ptr_coeff_2, p_v_out=p_v_out, p_u_out=p_u_out, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_vertices), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/model/common/tests/test_vertical.py b/model/common/tests/test_vertical.py new file mode 100644 index 000000000..4d6d71b9a --- /dev/null +++ b/model/common/tests/test_vertical.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import math + +import numpy as np +import pytest + +from icon4py.model.common.dimension import KDim +from icon4py.model.common.grid.vertical import VerticalModelParams + + +@pytest.mark.parametrize( + "max_h,damping,delta", + [(60000, 34000, 612), (12000, 10000, 100), (109050, 45000, 123)], +) +def test_nrdmax_calculation(max_h, damping, delta): + vct_a = np.arange(0, max_h, delta) + vct_a = vct_a[::-1] + vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) + assert vertical_params.index_of_damping_layer == vct_a.shape[0] - math.ceil(damping / delta) - 1 + + +@pytest.mark.datatest +def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): + a = grid_savepoint.vct_a() + nrdmax = grid_savepoint.nrdmax() + vertical_params = VerticalModelParams(rayleigh_damping_height=damping_height, vct_a=a) + assert nrdmax == vertical_params.index_of_damping_layer + a_array = np.asarray(a) + assert a_array[nrdmax] > damping_height + assert a_array[nrdmax + 1] < damping_height + + +@pytest.mark.datatest +def test_grid_size(grid_savepoint): + assert 65 == grid_savepoint.num(KDim) diff --git a/model/driver/.bumpversion.cfg b/model/driver/.bumpversion.cfg new file mode 100644 index 000000000..ae24af96d --- /dev/null +++ b/model/driver/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/driver/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/model/driver/.flake8 b/model/driver/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/driver/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/driver/.pre-commit-config.yaml b/model/driver/.pre-commit-config.yaml new file mode 100644 index 000000000..4ae0972f5 --- /dev/null +++ b/model/driver/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/driver/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/driver/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/driver/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/driver/.flake8, model/driver/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/driver/README.md b/model/driver/README.md new file mode 100644 index 000000000..aec461abb --- /dev/null +++ b/model/driver/README.md @@ -0,0 +1,36 @@ +# Dummy driver for Python ICON port + +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far the code mostly draws on serialized ICON data until we increasingly can initialize and run the model independently. + +It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. + +Currently, it does _no real timestepping_, instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. + +The code is meant to be changed and enlarged as we port new parts of the model. + +It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes are available. + +## Installation + +See the general instructions in the [README.md](../../README.md) in the base folder of the repository. + +## Usage + +```bash +export ICON4PY_ROOT= +dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +``` + +or if running in parallel + +```bash +mpirun -np 2 dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon +``` + +#### Remarks + +- First (required) arg is the folder where the serialized input data is stored. The input data is the same as is used in the unit tests. The path in the example is where the data is put when downloaded via the unit tests. + - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. +- parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. +- The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. +- The code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml new file mode 100644 index 000000000..b011e5f5b --- /dev/null +++ b/model/driver/pyproject.toml @@ -0,0 +1,126 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] + +[project] +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "pytz>=2023.2", + "icon4py-common>=0.0.5", + "icon4py-atmosphere-dycore>=0.0.5", + "icon4py-atmosphere-diffusion>=0.0.5" +] +description = "ICON model driver." +dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-driver" +readme = "README.md" +requires-python = ">=3.10" + +[project.scripts] +dycore_driver = "icon4py.model.driver.dycore_driver:main" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +branch = true +parallel = true +source_pkgs = ['driver'] + +[tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +multi_line_output = 3 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +use_parentheses = true + +[tool.mypy] +disallow_incomplete_defs = true +disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] +ignore_missing_imports = false +implicit_reexport = true +install_types = true +non_interactive = true +show_column_numbers = true +show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.driver.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.driver' = ['py.typed'] diff --git a/model/driver/requirements-dev.txt b/model/driver/requirements-dev.txt new file mode 100644 index 000000000..68e2b1610 --- /dev/null +++ b/model/driver/requirements-dev.txt @@ -0,0 +1,5 @@ +-r ../../base-requirements-dev.txt +-e ../common +-e ../atmosphere/diffusion +-e ../atmosphere/dycore +-e . diff --git a/model/driver/requirements.txt b/model/driver/requirements.txt new file mode 100644 index 000000000..18a81dc7c --- /dev/null +++ b/model/driver/requirements.txt @@ -0,0 +1,5 @@ +-r ../../base-requirements.txt +../common +../atmosphere/dycore +../atmosphere/diffusion +. diff --git a/model/driver/src/icon4py/model/driver/__init__.py b/model/driver/src/icon4py/model/driver/__init__.py new file mode 100644 index 000000000..dab708955 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/__init__.py @@ -0,0 +1,33 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py new file mode 100644 index 000000000..9103a2443 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -0,0 +1,295 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import logging +from datetime import datetime, timedelta +from pathlib import Path +from typing import Callable + +import click +import pytz +from devtools import Timer +from gt4py.next import Field, program +from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k +from icon4py.model.common.decomposition.decomposed import create_exchange +from icon4py.model.common.decomposition.parallel_setup import ( + ProcessProperties, + get_processor_properties, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils import serialbox_utils as sb +from icon4py.model.driver.icon_configuration import IconRunConfig, read_config +from icon4py.model.driver.io_utils import ( + SIMULATION_START_DATE, + configure_logging, + read_decomp_info, + read_geometry_fields, + read_icon_grid, + read_initial_state, + read_static_fields, +) + + +log = logging.getLogger(__name__) + + +# TODO (magdalena) to be removed once there is a proper time stepping +@program +def _copy_diagnostic_and_prognostics( + hdef_ic_new: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + div_ic_new: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + dwdx_new: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy_new: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + vn_new: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + w_new: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + exner_new: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + theta_v_new: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], +): + _identity_c_k(hdef_ic_new, out=hdef_ic) + _identity_c_k(div_ic_new, out=div_ic) + _identity_c_k(dwdx_new, out=dwdx) + _identity_c_k(dwdy_new, out=dwdy) + _identity_e_k(vn_new, out=vn) + _identity_c_k(w_new, out=w) + _identity_c_k(exner_new, out=exner) + _identity_c_k(theta_v_new, out=theta_v) + + +class DummyAtmoNonHydro: + def __init__(self, data_provider: sb.IconSerialDataProvider): + self.config = None + self.data_provider = data_provider + self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) + + def init(self, config): + self.config = config + + def _next_physics_date(self, dtime: float): + dynamics_dtime = dtime / self.config.n_substeps + self.simulation_date += timedelta(seconds=dynamics_dtime) + + def _dynamics_timestep(self, dtime): + """Show structure with this dummy fucntion called inside substepping loop.""" + self._next_physics_date(dtime) + + def do_dynamics_substepping( + self, + dtime, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + for _ in range(self.config.n_substeps): + self._dynamics_timestep(dtime) + sp = self.data_provider.from_savepoint_diffusion_init( + linit=False, date=self.simulation_date.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + ) + new_p = sp.construct_prognostics() + new_d = sp.construct_diagnostics_for_diffusion() + _copy_diagnostic_and_prognostics.with_backend(run_gtfn)( + new_d.hdef_ic, + diagnostic_state.hdef_ic, + new_d.div_ic, + diagnostic_state.div_ic, + new_d.dwdx, + diagnostic_state.dwdx, + new_d.dwdy, + diagnostic_state.dwdy, + new_p.vn, + prognostic_state.vn, + new_p.w, + prognostic_state.w, + new_p.exner_pressure, + prognostic_state.exner_pressure, + new_p.theta_v, + prognostic_state.theta_v, + offset_provider={}, + ) + + +class Timeloop: + @classmethod + def name(cls): + return cls.__name__ + + def __init__( + self, + config: IconRunConfig, + diffusion: Diffusion, + atmo_non_hydro: DummyAtmoNonHydro, + ): + self.config = config + self.diffusion = diffusion + self.atmo_non_hydro = atmo_non_hydro + + def _full_name(self, func: Callable): + return ":".join((self.__class__.__name__, func.__name__)) + + def _timestep( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + + self.atmo_non_hydro.do_dynamics_substepping( + self.config.dtime, diagnostic_state, prognostic_state + ) + self.diffusion.run( + diagnostic_state, + prognostic_state, + self.config.dtime, + ) + + def __call__( + self, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + ): + log.info( + f"starting time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + log.info("running initial step to diffuse fields before timeloop starts") + self.diffusion.initial_run( + diagnostic_state, + prognostic_state, + self.config.dtime, + ) + log.info( + f"starting real time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + timer = Timer(self._full_name(self._timestep)) + for t in range(self.config.n_time_steps): + log.info(f"run timestep : {t}") + timer.start() + self._timestep(diagnostic_state, prognostic_state) + timer.capture() + timer.summary(True) + + +def initialize(n_time_steps, file_path: Path, props: ProcessProperties): + """ + Inititalize the driver run. + + "reads" in + - configuration + + - grid information + + - (serialized) input fields, initial + + Returns: + tl: configured timeloop, + prognostic_state: initial state fro prognostic and diagnostic variables + diagnostic_state: + """ + log.info("initialize parallel runtime") + experiment_name = "mch_ch_r04b09_dsl" + log.info(f"reading configuration: experiment {experiment_name}") + config = read_config(experiment_name, n_time_steps=n_time_steps) + + decomp_info = read_decomp_info(file_path, props) + + log.info(f"initializing the grid from '{file_path}'") + icon_grid = read_icon_grid(file_path, rank=props.rank) + log.info(f"reading input fields from '{file_path}'") + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( + file_path, rank=props.rank + ) + (metric_state, interpolation_state) = read_static_fields(file_path) + + log.info("initializing diffusion") + diffusion_params = DiffusionParams(config.diffusion_config) + exchange = create_exchange(props, decomp_info) + diffusion = Diffusion(exchange) + diffusion.init( + icon_grid, + config.diffusion_config, + diffusion_params, + vertical_geometry, + metric_state, + interpolation_state, + edge_geometry, + cell_geometry, + ) + + data_provider, diagnostic_state, prognostic_state = read_initial_state( + file_path, rank=props.rank + ) + + atmo_non_hydro = DummyAtmoNonHydro(data_provider) + atmo_non_hydro.init(config=config.dycore_config) + + tl = Timeloop( + config=config.run_config, + diffusion=diffusion, + atmo_non_hydro=atmo_non_hydro, + ) + return tl, diagnostic_state, prognostic_state + + +@click.command() +@click.argument("input_path") +@click.option("--run_path", default="", help="folder for output") +@click.option("--n_steps", default=5, help="number of time steps to run, max 5 is supported") +@click.option("--mpi", default=False, help="whether or not you are running with mpi") +def main(input_path, run_path, n_steps, mpi): + """ + Run the driver. + + usage: python driver/dycore_driver.py ../../tests/ser_icondata/mch_ch_r04b09_dsl/ser_data + + steps: + 1. initialize model: + + a) load config + + b) initialize grid + + c) initialize/configure components ie "granules" + + d) setup the time loop + + 2. run time loop + + """ + start_time = datetime.now().astimezone(pytz.UTC) + parallel_props = get_processor_properties(with_mpi=mpi) + configure_logging(run_path, start_time, parallel_props) + log.info(f"Starting ICON dycore run: {datetime.isoformat(start_time)}") + log.info(f"input args: input_path={input_path}, n_time_steps={n_steps}") + timeloop, diagnostic_state, prognostic_state = initialize( + n_steps, Path(input_path), parallel_props + ) + log.info("dycore configuring: DONE") + log.info("timeloop: START") + + timeloop(diagnostic_state, prognostic_state) + + log.info("timeloop: DONE") + + +if __name__ == "__main__": + main() diff --git a/model/driver/src/icon4py/model/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py new file mode 100644 index 000000000..187c4f4f4 --- /dev/null +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -0,0 +1,85 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import dataclass +from typing import Optional + +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType + + +n_substeps_reduced = 2 + + +@dataclass +class IconRunConfig: + n_time_steps: int = 5 + dtime: float = 600.0 + + +@dataclass +class AtmoNonHydroConfig: + n_substeps: int = 5 + + +@dataclass +class IconConfig: + run_config: IconRunConfig + diffusion_config: DiffusionConfig + dycore_config: AtmoNonHydroConfig + + +def read_config(experiment: Optional[str], n_time_steps: int) -> IconConfig: + def _default_run_config(n_steps: int): + if n_steps > 5: + raise NotImplementedError("only five dummy timesteps available") + return IconRunConfig(n_time_steps=n_steps) + + def mch_ch_r04b09_diffusion_config(): + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + n_substeps=n_substeps_reduced, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + ) + + def _default_diffusion_config(): + return DiffusionConfig() + + def _default_config(n_steps): + run_config = _default_run_config(n_steps) + return run_config, _default_diffusion_config(), AtmoNonHydroConfig() + + def _mch_ch_r04b09_config(n_steps): + return ( + IconRunConfig(n_time_steps=n_steps, dtime=10.0), + mch_ch_r04b09_diffusion_config(), + AtmoNonHydroConfig(), + ) + + if experiment == "mch_ch_r04b09_dsl": + (model_run_config, diffusion_config, dycore_config) = _mch_ch_r04b09_config(n_time_steps) + else: + (model_run_config, diffusion_config, dycore_config) = _default_config(n_time_steps) + return IconConfig( + run_config=model_run_config, + diffusion_config=diffusion_config, + dycore_config=dycore_config, + ) diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py new file mode 100644 index 000000000..ae14933bb --- /dev/null +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -0,0 +1,187 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import logging +from datetime import datetime +from enum import Enum +from pathlib import Path + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.common.decomposition.decomposed import DecompositionInfo +from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.icon_grid import IconGrid +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils import serialbox_utils as sb + + +SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." + +SIMULATION_START_DATE = "2021-06-20T12:00:10.000" +log = logging.getLogger(__name__) + + +class SerializationType(str, Enum): + SB = "serialbox" + NC = "netcdf" + + +def read_icon_grid( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> IconGrid: + """ + Read icon grid. + + Args: + path: path where to find the input data + ser_type: type of input data. Currently only 'sb (serialbox)' is supported. It reads + from ppser serialized test data + Returns: IconGrid parsed from a given input type. + """ + if ser_type == SerializationType.SB: + return ( + sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) + .from_savepoint_grid() + .construct_icon_grid() + ) + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_initial_state( + gridfile_path: Path, rank=0 +) -> tuple[sb.IconSerialDataProvider, DiffusionDiagnosticState, PrognosticState]: + """ + Read prognostic and diagnostic state from serialized data. + + Args: + gridfile_path: path the serialized input data + + Returns: a tuple containing the data_provider, the initial diagnostic and prognostic state. + The data_provider is returned such that further timesteps of diagnostics and prognostics + can be read from within the dummy timeloop + + """ + data_provider = sb.IconSerialDataProvider( + "icon_pydycore", str(gridfile_path), False, mpi_rank=rank + ) + init_savepoint = data_provider.from_savepoint_diffusion_init( + linit=True, date=SIMULATION_START_DATE + ) + prognostic_state = init_savepoint.construct_prognostics() + diagnostic_state = init_savepoint.construct_diagnostics_for_diffusion() + return data_provider, diagnostic_state, prognostic_state + + +def read_geometry_fields( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> tuple[EdgeParams, CellParams, VerticalModelParams]: + """ + Read fields containing grid properties. + + Args: + path: path to the serialized input data + ser_type: (optional) defaults to SB=serialbox, type of input data to be read + + Returns: a tuple containing fields describing edges, cells, vertical properties of the model + the data is originally obtained from the grid file (horizontal fields) or some special input files. + """ + if ser_type == SerializationType.SB: + sp = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank + ).from_savepoint_grid() + edge_geometry = sp.construct_edge_geometry() + cell_geometry = sp.construct_cell_geometry() + vertical_geometry = VerticalModelParams(vct_a=sp.vct_a(), rayleigh_damping_height=12500) + return edge_geometry, cell_geometry, vertical_geometry + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_decomp_info( + path: Path, + procs_props: ProcessProperties, + ser_type=SerializationType.SB, +) -> DecompositionInfo: + if ser_type == SerializationType.SB: + sp = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), True, procs_props.rank + ) + return sp.from_savepoint_grid().construct_decomposition_info() + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def read_static_fields( + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB +) -> tuple[DiffusionMetricState, DiffusionInterpolationState]: + """ + Read fields for metric and interpolation state. + + Args: + path: path to the serialized input data + rank: mpi rank, defaults to 0 for serial run + ser_type: (optional) defaults to SB=serialbox, type of input data to be read + + Returns: + a tuple containing the metric_state and interpolation state, + the fields are precalculated in the icon setup. + + """ + if ser_type == SerializationType.SB: + dataprovider = sb.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank + ) + interpolation_state = ( + dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() + ) + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + return metric_state, interpolation_state + else: + raise NotImplementedError(SB_ONLY_MSG) + + +def configure_logging(run_path: str, start_time, processor_procs: ProcessProperties = None) -> None: + """ + Configure logging. + + Log output is sent to console and to a file. + + Args: + run_path: path to the output folder where the logfile should be stored + start_time: start time of the model run + + """ + run_dir = Path(run_path).absolute() if run_path else Path(__file__).absolute().parent + run_dir.mkdir(exist_ok=True) + logfile = run_dir.joinpath(f"dummy_dycore_driver_{datetime.isoformat(start_time)}.log") + logfile.touch(exist_ok=True) + logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s %(filename)-20s (%(lineno)-4d) : %(funcName)-20s: %(levelname)-8s %(message)s", + filemode="w", + filename=logfile, + ) + console_handler = logging.StreamHandler() + console_handler.addFilter(ParallelLogger(processor_procs)) + + log_format = "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" + formatter = logging.Formatter(fmt=log_format, style="{", defaults={"rank": None}) + console_handler.setFormatter(formatter) + console_handler.setLevel(logging.DEBUG) + logging.getLogger("").addHandler(console_handler) diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py new file mode 100644 index 000000000..d9ad439dc --- /dev/null +++ b/model/driver/tests/conftest.py @@ -0,0 +1,24 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 + datapath, + download_ser_data, + processor_props, + ranked_data_path, +) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/driver/tests/test_io_utils.py b/model/driver/tests/test_io_utils.py new file mode 100644 index 000000000..de51e2af1 --- /dev/null +++ b/model/driver/tests/test_io_utils.py @@ -0,0 +1,108 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest + +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.driver.io_utils import ( + SerializationType, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) + + +@pytest.mark.datatest +@pytest.mark.parametrize("read_fun", (read_geometry_fields, read_static_fields, read_icon_grid)) +def test_read_geometry_fields_not_implemented_type(read_fun, datapath): + with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): + read_fun(path=datapath, ser_type=SerializationType.NC) + + +def assert_grid_size_and_connectivities(grid): + assert grid.num_edges() == 31558 + assert grid.num_cells() == 20896 + assert grid.num_vertices() == 10663 + assert grid.get_e2v_connectivity() + assert grid.get_v2e_connectivity() + assert grid.get_c2e_connectivity() + assert grid.get_e2c_connectivity() + assert grid.get_e2c2v_connectivity() + assert grid.get_c2e2c_connectivity() + assert grid.get_e2ecv_connectivity() + + +@pytest.mark.datatest +def test_read_icon_grid_for_type_sb(datapath): + grid = read_icon_grid(datapath, ser_type=SerializationType.SB) + assert_grid_size_and_connectivities(grid) + + +@pytest.mark.datatest +def test_read_static_fields_for_type_sb(datapath): + metric_state, interpolation_state = read_static_fields(datapath, ser_type=SerializationType.SB) + assert_metric_state_fields(metric_state) + assert_interpolation_state_fields(interpolation_state) + + +@pytest.mark.datatest +def test_read_geometry_fields_for_type_sb(datapath): + edge_geometry, cell_geometry, vertical_geometry = read_geometry_fields( + datapath, ser_type=SerializationType.SB + ) + assert_edge_geometry_fields(edge_geometry) + assert_cell_geometry_fields(cell_geometry) + assert_vertical_params(vertical_geometry) + + +def assert_vertical_params(vertical_geometry: VerticalModelParams): + assert vertical_geometry.physical_heights + assert vertical_geometry.index_of_damping_layer > 0 + assert vertical_geometry.rayleigh_damping_height > 0 + + +def assert_cell_geometry_fields(cell_geometry: CellParams): + assert cell_geometry.area + + +def assert_edge_geometry_fields(edge_geometry: EdgeParams): + assert edge_geometry.edge_areas + assert edge_geometry.primal_normal_vert + assert edge_geometry.inverse_primal_edge_lengths + assert edge_geometry.tangent_orientation + assert edge_geometry.inverse_dual_edge_lengths + assert edge_geometry.dual_normal_vert + + +def assert_metric_state_fields(metric_state): + assert metric_state.wgtfac_c + assert metric_state.zd_intcoef + assert metric_state.zd_diffcoef + assert metric_state.theta_ref_mc + assert metric_state.mask_hdiff + assert metric_state.zd_vertoffset + + +def assert_interpolation_state_fields(interpolation_state): + assert interpolation_state.geofac_n2s + assert interpolation_state.e_bln_c_s + assert interpolation_state.nudgecoeff_e + assert interpolation_state.geofac_n2s_nbh + assert interpolation_state.geofac_div + assert interpolation_state.geofac_grg_y + assert interpolation_state.geofac_grg_x + assert interpolation_state.rbf_coeff_2 + assert interpolation_state.rbf_coeff_1 + assert interpolation_state.geofac_n2s_c diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index 1b7e16d5a..206e27f11 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -1,3 +1,5 @@ -r ../base-requirements-dev.txt -e ./atmosphere/dycore +-e ./atmosphere/diffusion -e ./common +-e ./driver diff --git a/model/requirements.txt b/model/requirements.txt index 2045fe905..519a20834 100644 --- a/model/requirements.txt +++ b/model/requirements.txt @@ -1,3 +1,5 @@ -r ../base-requirements.txt ./atmosphere/dycore +./atmosphere/diffusion ./common +./driver diff --git a/model/tox.ini b/model/tox.ini index 27fb18d61..9ef46bd3c 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src common/src driver/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html @@ -31,3 +31,4 @@ setenv = skip_install = true commands = commands_post = + diff --git a/requirements-dev.txt b/requirements-dev.txt index 03fc169c9..ca807eac7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,6 +3,8 @@ # icon4py model -e ./model/atmosphere/dycore -e ./model/common +-e ./model/atmosphere/diffusion +-e ./model/driver # icon4pytools -e ./tools diff --git a/requirements.txt b/requirements.txt index 36a728881..14c1cc762 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,9 @@ # icon4py model ./model/atmosphere/dycore +./model/atmosphere/diffusion ./model/common +./model/driver # icon4pytools ./tools diff --git a/tools/README.md b/tools/README.md index fdb881588..41f040e0d 100644 --- a/tools/README.md +++ b/tools/README.md @@ -232,3 +232,39 @@ OUTPUT_FILEPATH A path to the output Fortran source file to be generated. ``` **Note:** The output of f2ser still has to be preprocessed using `pp_ser.py`, which then yields a compilable unit. The serialised files will have `f2ser` as their prefix in the default folder location of the experiment. + +## `py2f` + +Python utility for generating a C library and Fortran interface to call Python icon4py modules. The library [embeds python via CFFI ](https://cffi.readthedocs.io/en/latest/embedding.) + +This is **highly experimental** and has not been tested from within Fortran code! + +### usage + +`py2fgen`: Generates a C header file and a Fortran interface and compiles python functions into a C library embedding python. The functions need to be decorated with `CffiMethod.register` and have a signature with scalar arguments or `GT4Py` fields, for example: + +``` +@CffiMethod.register +def foo(i:int, param:float, field1: Field[[VertexDim, KDim], float], field2: Field[CellDim, KDim], float]) +``` + +see `src/icon4pytools/py2f/wrappers` for examples. + +```bash +py2fgen icon4pytools.py2f.wrappers.diffusion_wrapper py2f_build +``` + +where the first argument is the python module to parse and the second a build directory. The call above will generate the following in `py2f_build`: + +```bash + ls py2f_build/ +total 204K +drwxrwxr-x 2 magdalena magdalena 4.0K Aug 24 17:01 . +drwxrwxr-x 9 magdalena magdalena 4.0K Aug 24 17:01 .. +-rw-rw-r-- 1 magdalena magdalena 74K Aug 24 16:58 diffusion_wrapper.c +-rw-rw-r-- 1 magdalena magdalena 3.5K Aug 24 17:01 diffusion_wrapper.f90 +-rw-rw-r-- 1 magdalena magdalena 955 Aug 24 17:01 diffusion_wrapper.h +-rw-rw-r-- 1 magdalena magdalena 58K Aug 24 17:01 diffusion_wrapper.o +-rwxrwxr-x 1 magdalena magdalena 49K Aug 24 17:01 libdiffusion_wrapper.so + +``` diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 0c828866a..9df4270d6 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -19,11 +19,13 @@ classifiers = [ 'Topic :: Scientific/Engineering :: Physics' ] dependencies = [ + 'icon4py-atmosphere-diffusion', 'icon4py-atmosphere-dycore', 'gt4py>=1.0.1', - 'icon4py_common', + 'icon4py-common', 'tabulate>=0.8.9', - 'fprettify>=0.3.7' + 'fprettify>=0.3.7', + 'cffi>=1.5' ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] @@ -36,6 +38,7 @@ requires-python = '>=3.10' f2ser = 'icon4pytools.f2ser.cli:main' icon4pygen = 'icon4pytools.icon4pygen.cli:main' icon_liskov = 'icon4pytools.liskov.cli:main' +py2fgen = 'icon4pytools.py2f.py2fgen:main' [project.urls] repository = 'https://github.com/C2SM/icon4py' @@ -108,7 +111,12 @@ use_parentheses = true [tool.mypy] disallow_incomplete_defs = true disallow_untyped_defs = true -exclude = ['^tests/f2ser/*.py', '^tests/icon4pygen/*.py', '^tests/liskov/*.py'] +exclude = [ + '^tests/f2ser/*.py', + '^tests/icon4pygen/*.py', + '^tests/liskov/*.py', + '^tests/py2f/*.py' +] ignore_missing_imports = false implicit_reexport = true install_types = true diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index f3d56fb5c..d4652678b 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,4 +1,5 @@ -r ../base-requirements-dev.txt -e ../model/atmosphere/dycore +-e ../model/atmosphere/diffusion -e ../model/common -e . diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 53b815173..6edcc302a 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -32,7 +32,11 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atmosphere.dycore"] + _STENCIL_PACKAGES = [ + "atmosphere.dycore", + "atmosphere.diffusion.stencils", + "common.interpolation.stencils", + ] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed diff --git a/tools/src/icon4pytools/py2f/__init__.py b/tools/src/icon4pytools/py2f/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/tools/src/icon4pytools/py2f/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py new file mode 100644 index 000000000..c25183132 --- /dev/null +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -0,0 +1,177 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import functools +import inspect +from collections import OrderedDict +from importlib.resources import files +from types import MappingProxyType +from typing import Any + +import cffi +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4pytools.py2f.typing_utils import parse_annotation + + +FFI_DEF_EXTERN_DECORATOR = "@ffi.def_extern()" + +CFFI_GEN_DECORATOR = "@CffiMethod.register" + + +class CffiMethod: + _registry = {} + + @classmethod + def register(cls, func): + cls._registry.setdefault(func.__module__, []).append(func.__name__) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + + @classmethod + def get(cls, name: str): + return cls._registry[name] + + +def generate_and_compile_cffi_plugin( + plugin_name: str, c_header: str, module_name: str, build_path="." +): + """ + Create C shared library. + + Create a linkable C library and F90 interface for the functions in the python module + {module_name} that are decorated with '@CffiMethod.register'. + + Args: + plugin_name: name of the plugin, a linkable C library with the name + 'lib{plugin_name}.so' will be created in the {build_path} folder' + c_header: C type header signature for the python functions. + module_name: python module name that contains python functions corresponding + to the signature in the '{c_header}' string, these functions must be decorated + with @CffiMethod.register and the file must contain the import + build_path: *optional* path to build directory + + """ + module_split = module_name.split(".") + python_src_file = "/".join(module_split[1:]) + ".py" + python_package = module_split[0] + + c_header_file = plugin_name + ".h" + with open("/".join([build_path, c_header_file]), "w") as f: + f.write(c_header) + + builder = cffi.FFI() + + builder.embedding_api(c_header) + builder.set_source(plugin_name, f'#include "{c_header_file}"') + + module = files(python_package).joinpath(python_src_file).read_text() + + module = f"from {plugin_name} import ffi\n{module}".replace( + CFFI_GEN_DECORATOR, FFI_DEF_EXTERN_DECORATOR + ) + + builder.embedding_init_code(module) + builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) + + +class UnknownDimensionException(Exception): + """Raised if a Dimension is unknown to the interface generation.""" + + pass + + +def to_fields(dim_sizes: dict[Dimension, int]): + """ + Pack/Unpack Fortran 2d arrays to numpy arrays with using CFFI frombuffer. + + Args: + dim_sizes: dictionary containing the sizes of the dimension. + + #TODO (magdalena) handle dimension sizes in a better way? + """ + ffi = cffi.FFI() + dim_sizes = dim_sizes + + def _dim_sizes(dims: list[Dimension]) -> tuple[int, int]: + """Extract the size of dimension from a dictionary.""" + v_size = None + h_size = None + for d in dims: + if d not in dim_sizes.keys(): + raise UnknownDimensionException( + f"size of dimension '{d}' not defined in '{dim_sizes.keys()}'" + ) + if d.kind == DimensionKind.VERTICAL or d.kind == DimensionKind.LOCAL: + v_size = dim_sizes[d] + elif d.kind == DimensionKind.HORIZONTAL: + h_size = dim_sizes[d] + return h_size, v_size + + def _unpack(ptr, size_h, size_v, dtype) -> np.ndarray: + """ + Unpack a 2d c/fortran field into a numpy array. + + :param dtype: expected type of the fields + :param ptr: c_pointer to the field + :param size_h: length of horizontal dimension + :param size_v: length of vertical dimension + :return: a numpy array with shape=(size_h, size_v) + and dtype = ctype of the pointer + """ + shape = (size_h, size_v) + length = np.prod(shape) + c_type = ffi.getctype(ffi.typeof(ptr).item) + # TODO (magdalena) fix dtype handling use SCALARTYPE? + mem_size = ffi.sizeof(c_type) + mem_size = np.dtype(c_type).itemsize + ar = np.frombuffer( + ffi.buffer(ptr, length * mem_size), + dtype=np.dtype(c_type), + count=-1, + offset=0, + ).reshape(shape) + return ar + + def _to_fields_decorator(func): + @functools.wraps(func) + def _wrapper( + *args, **kwargs + ): # these are the args of the decorated function ie to_fields(func(*args, **kwargs)) + signature = inspect.signature(func) + parameters: MappingProxyType[str, inspect.Parameter] = signature.parameters + arguments: OrderedDict[str, Any] = signature.bind(*args, **kwargs).arguments + f_args = [] + for name, argument in arguments.items(): + ar = _transform_arg(argument, name, parameters) + f_args.append(ar) + return func(*f_args, **kwargs) + + def _transform_arg(argument, name, parameters): + dims, dtype = parse_annotation(parameters[name].annotation) + if dims: + (size_h, size_v) = _dim_sizes(dims) + ar = _unpack(argument, size_h, size_v, dtype) + ar = np_as_located_field(*dims)(ar) + else: + ar = argument + return ar + + return _wrapper + + return _to_fields_decorator diff --git a/tools/src/icon4pytools/py2f/codegen.py b/tools/src/icon4pytools/py2f/codegen.py new file mode 100644 index 000000000..502b20492 --- /dev/null +++ b/tools/src/icon4pytools/py2f/codegen.py @@ -0,0 +1,145 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Sequence + +from gt4py.eve import Node, codegen +from gt4py.eve.codegen import JinjaTemplate as as_jinja +from gt4py.eve.codegen import TemplatedGenerator +from gt4py.next.type_system.type_specifications import ScalarKind + +from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( + BUILTIN_TO_CPP_TYPE, + BUILTIN_TO_ISO_C_TYPE, +) +from icon4pytools.icon4pygen.bindings.utils import write_string + + +class DimensionType(Node): + name: str + length: int + + +class FuncParameter(Node): + name: str + d_type: ScalarKind + dimensions: Sequence[DimensionType] + + +class Func(Node): + name: str + args: Sequence[FuncParameter] + + +class CffiPlugin(Node): + name: str + functions: Sequence[Func] + + +def to_c_type(scalar_type: ScalarKind) -> str: + return BUILTIN_TO_CPP_TYPE[scalar_type] + + +def to_f_type(scalar_type: ScalarKind) -> str: + return BUILTIN_TO_ISO_C_TYPE[scalar_type] + + +def as_f90_value(param: FuncParameter) -> str: + """ + If param is a scalar type (dimension=0) then return the F90 'value' keyword. + + Used for F90 generation only. + """ + return "value, " if len(param.dimensions) == 0 else "" + + +def as_field(param: FuncParameter, language: str) -> str: + size = len(param.dimensions) + if size == 0: + return "" + if "C" == language: + return "*" + else: + dims = ",".join(map(lambda x: ":", range(size))) + return f"({dims})" + + +class CHeaderGenerator(TemplatedGenerator): + CffiPlugin = as_jinja("""{% for func in functions: %}{{func}}\n{% endfor %}""") + + Func = as_jinja("""extern void {{name}}({{", ".join(args)}});""") + + def visit_FuncParameter(self, param: FuncParameter): + return self.generic_visit( + param, + rendered_type=to_c_type(param.d_type), + dim=as_field(param, "C"), + ) + + FuncParameter = as_jinja("""{{rendered_type}}{{dim}} {{name}}""") + + +class F90InterfaceGenerator(TemplatedGenerator): + CffiPlugin = as_jinja( + """ + module {{name}} + use, intrinsic:: iso_c_binding + implicit none + + public + interface + {% for func in functions: %}\ + {{func}}\ + {% endfor %}\ + end interface + end module + """ + ) + + def visit_Func(self, func: Func): + arg_names = ", ".join(map(lambda x: x.name, func.args)) + return self.generic_visit(func, param_names=arg_names) + + Func = as_jinja( + """subroutine {{name}}({{param_names}}) bind(c, name='{{name}}') + use iso_c_binding + {% for arg in args: %}\ + {{arg}}\ + {% endfor %}\ + end subroutine {{name}} + """ + ) + + def visit_FuncParameter(self, param: FuncParameter, param_names=""): + # kw-arg param_names needs to be present because it is present up the tree + return self.generic_visit( + param, + value=as_f90_value(param), + rendered_type=to_f_type(param.d_type), + dim=as_field(param, "F"), + ) + + FuncParameter = as_jinja( + """{{rendered_type}}, {{value}} intent(inout):: {{name}}{{dim}} + """ + ) + + +def generate_c_header(plugin: CffiPlugin) -> str: + generated_code = CHeaderGenerator.apply(plugin) + return codegen.format_source("cpp", generated_code, style="LLVM") + + +def generate_and_write_f90_interface(build_path: str, plugin: CffiPlugin): + generated_code = F90InterfaceGenerator.apply(plugin) + write_string(generated_code, build_path, f"{plugin.name}.f90") diff --git a/tools/src/icon4pytools/py2f/parsing.py b/tools/src/icon4pytools/py2f/parsing.py new file mode 100644 index 000000000..847285615 --- /dev/null +++ b/tools/src/icon4pytools/py2f/parsing.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import importlib +from inspect import signature, unwrap + +from icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter +from icon4pytools.py2f.typing_utils import parse_annotation + + +def parse_functions_from_module(module_name: str) -> CffiPlugin: + module = importlib.import_module(module_name) + func_names = CffiMethod.get(module_name) + funcs = [_parse_function(module, fn) for fn in func_names] + plugin_name = module_name.split(".")[-1] + return CffiPlugin(name=plugin_name, functions=funcs) + + +def _parse_function(module, s): + func = unwrap(getattr(module, s)) + params = [ + _parse_params(signature(func, follow_wrapped=False).parameters, p) + for p in (signature(func).parameters) + ] + return Func(name=s, args=params) + + +def _parse_params(params, s): + annotation = params[s].annotation + dims, dtype = parse_annotation(annotation) + dim_types = [DimensionType(name=d.value, length=10) for d in dims] + return FuncParameter(name=s, d_type=dtype, dimensions=dim_types) diff --git a/tools/src/icon4pytools/py2f/py2fgen.py b/tools/src/icon4pytools/py2f/py2fgen.py new file mode 100644 index 000000000..7bdb799d1 --- /dev/null +++ b/tools/src/icon4pytools/py2f/py2fgen.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pathlib + +import click + +from icon4pytools.py2f.cffi_utils import generate_and_compile_cffi_plugin +from icon4pytools.py2f.codegen import generate_and_write_f90_interface, generate_c_header +from icon4pytools.py2f.parsing import parse_functions_from_module + + +@click.command( + "py2fgen", +) +@click.argument("module", type=str) +@click.argument( + "build_path", + type=click.Path(dir_okay=True, resolve_path=True, path_type=pathlib.Path), + default=".", +) +def main(module: str, build_path: pathlib.Path) -> None: + """ + Generate C and F90 wrappers and C library for embedding the python MODULE in C and Fortran. + + Args: + - module: name of the python module containing the methods to be embedded. Those + methods have to be decorated with CffiMethod.register + + - build_path: directory where the generated code and compiled libraries are to be found. + """ + module_name = module + build_path.mkdir(exist_ok=True, parents=True) + plugin = parse_functions_from_module(module_name) + c_header = generate_c_header(plugin) + generate_and_compile_cffi_plugin(plugin.name, c_header, module_name, str(build_path)) + generate_and_write_f90_interface(build_path, plugin) + + +if __name__ == "__main__": + main() diff --git a/tools/src/icon4pytools/py2f/typing_utils.py b/tools/src/icon4pytools/py2f/typing_utils.py new file mode 100644 index 000000000..72ba5fa28 --- /dev/null +++ b/tools/src/icon4pytools/py2f/typing_utils.py @@ -0,0 +1,27 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Dimension +from gt4py.next.type_system.type_specifications import ScalarType +from gt4py.next.type_system.type_translation import from_type_hint + + +def parse_annotation(annotation) -> tuple[list[Dimension], ScalarType]: + type_spec = from_type_hint(annotation) + if isinstance(type_spec, ScalarType): + dtype = type_spec.kind + dims = [] + else: + dtype = type_spec.dtype.kind + dims = type_spec.dims + return dims, dtype diff --git a/tools/src/icon4pytools/py2f/wrappers/__init__.py b/tools/src/icon4pytools/py2f/wrappers/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/tools/src/icon4pytools/py2f/wrappers/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py new file mode 100644 index 000000000..89c2dbb59 --- /dev/null +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -0,0 +1,265 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# flake8: noqa +""" +Wrapper module for diffusion granule. + +Module contains a diffusion_init and diffusion_run function that follow the architecture of +Fortran granule interfaces: +- all arguments needed from external sources are passed. +- passing of scalar types or fields of simple types + + +""" +import numpy as np +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionConfig, + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + DiffusionParams, + DiffusionType, + PrognosticState, +) +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CECDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid +from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams + +from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields + + +# TODO (magdalena) Revise interface architecture with Fortran granules: +# The module variable to match the Fortran interface: where only fields are passed. +# We should rather instantiate the object init and return it. +diffusion: Diffusion() + +nproma = 50000 +field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} + + +@to_fields(dim_sizes=field_sizes) +@CffiMethod.register +def diffusion_init( + nproma: int, + nlev: int, + n_shift_total: int, + nrdmax: float, + n_dyn_substeps: int, + nudge_max_coeff: float, + denom_diffu_v: float, + hdiff_smag_fac: float, + hdiff_order: int, + hdiff_efdt_ratio: float, + itype_sher: int, + itype_vn_diffu: int, + itype_t_diffu: int, + lhdiff_rcf: bool, + lhdiff_w: bool, + lhdiff_temp: bool, + l_limited_area: bool, + l_zdiffu_t: bool, + lsmag_3d: bool, + vct_a: Field[[KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + nudgecoeff_e: Field[[EdgeDim], float], + zd_intcoef: Field[ + [CellDim, C2E2CDim, KDim], float + ], # special DSL field: zd_intcoef_dsl in mo_vertical_grid.f90 + zd_diffcoef: Field[ + [CellDim, KDim], float + ], # special DSL field mask instead of list: zd_diffcoef_dsl in mo_vertical_grid.f90 + wgtfac_c: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + edges_vertex_idx: Field[[EdgeDim, E2VDim], int32], + edges_cell_idx: Field[[EdgeDim, E2CDim], int32], + edges_tangent_orientation: Field[[EdgeDim], float], + edges_primal_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_primal_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_inv_vert_vert_lengths: Field[[EdgeDim], float], + edges_inv_primal_edge_length: Field[[EdgeDim], float], + edges_inv_dual_edge_length: Field[[EdgeDim], float], + edges_area_edge: Field[[EdgeDim], float], + cells_neighbor_idx: Field[[CellDim, C2E2CDim], int32], + cells_edge_idx: Field[[CellDim, C2EDim], int32], + cells_area: Field[[CellDim], float], + # dsl specific additional args + mask_hdiff: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[ + [CECDim], int32 + ], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 + rbf_coeff_1: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + rbf_coeff_2: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + verts_edge_idx: Field[[VertexDim, V2EDim], int32], # -> mo_intp_rbf_rbf_vec_interpol_vertex +): + """ + Instantiate and Initialize the diffusion object. + + should only accept simple fields as arguments for compatibility with the standalone + Fortran ICON Diffusion component (aka Diffusion granule) + + """ + if diffusion.initialized: + raise DuplicateInitializationException("Diffusion has already been already initialized") + + horizontal_size = HorizontalGridSize(num_cells=nproma, num_vertices=nproma, num_edges=nproma) + vertical_size = VerticalGridSize(num_lev=nlev) + mesh_config = GridConfig( + horizontal_config=horizontal_size, + vertical_config=vertical_size, + limited_area=l_limited_area, + n_shift_total=n_shift_total, + ) + # we need the start, end indices in order for the grid to be functional those are not passed + # to init in the Fortran diffusion granule since they are hidden away in the get_indices_[c,e,v] + # diffusion_run does take p_patch in order to pass it on to other subroutines (interpolation, get_indices... + + c2e2c0 = np.column_stack( + ( + (np.asarray(cells_neighbor_idx)), + (np.asarray(range(np.asarray(cells_neighbor_idx).shape[0]))), + ) + ) + e2c2v = np.asarray(edges_vertex_idx) + e2v = e2c2v[:, 0:2] + connectivities = { + E2CDim: np.asarray(edges_cell_idx), + E2C2VDim: e2c2v, + E2VDim: e2v, + C2EDim: np.asarray(cells_edge_idx), + C2E2CDim: (np.asarray(cells_neighbor_idx)), + C2E2CODim: c2e2c0, + V2EDim: np.asarray(verts_edge_idx), + } + # TODO (Magdalena) we need start_index, end_index: those are not passed in the fortran granules, + # because they are used through get_indices only + grid = IconGrid().with_config(mesh_config).with_connectivities(connectivities) + edge_params = EdgeParams( + tangent_orientation=edges_tangent_orientation, + inverse_primal_edge_lengths=edges_inv_primal_edge_length, + dual_normal_vert_x=edges_dual_normal_vert_1, + dual_normal_vert_y=edges_dual_normal_vert_2, + inverse_dual_edge_lengths=edges_inv_dual_edge_length, + inverse_vertex_vertex_lengths=edges_inv_vert_vert_lengths, + primal_normal_vert_x=edges_primal_normal_vert_1, + primal_normal_vert_y=edges_primal_normal_vert_2, + edge_areas=edges_area_edge, + ) + cell_params = CellParams(cells_area) + config: DiffusionConfig = DiffusionConfig( + diffusion_type=DiffusionType(hdiff_order), + hdiff_w=lhdiff_w, + hdiff_temp=lhdiff_temp, + type_vn_diffu=itype_vn_diffu, + smag_3d=lsmag_3d, + type_t_diffu=itype_t_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + hdiff_w_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=hdiff_smag_fac, + n_substeps=n_dyn_substeps, + zdiffu_t=l_zdiffu_t, + hdiff_rcf=lhdiff_rcf, + velocity_boundary_diffusion_denom=denom_diffu_v, + max_nudging_coeff=nudge_max_coeff, + ) + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) + + derived_diffusion_params = DiffusionParams(config) + metric_state = DiffusionMetricState( + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + mask_hdiff=mask_hdiff, + zd_vertoffset=zd_vertoffset, + zd_diffcoef=zd_diffcoef, + zd_intcoef=zd_intcoef, + ) + interpolation_state = DiffusionInterpolationState( + e_bln_c_s, + rbf_coeff_1, + rbf_coeff_2, + geofac_div, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + nudgecoeff_e, + ) + + diffusion.init( + grid=grid, + config=config, + params=derived_diffusion_params, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + +@CffiMethod.register +def diffusion_run( + dtime: float, + linit: bool, + vn: Field[[EdgeDim, KDim], float], + w: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], +): + diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + prognostic_state = PrognosticState( + w=w, + vn=vn, + exner_pressure=exner, + theta_v=theta_v, + ) + if linit: + diffusion.initial_run( + diagnostic_state, + prognostic_state, + dtime, + ) + else: + diffusion.run(diagnostic_state, prognostic_state, dtime) + + +class DuplicateInitializationException(Exception): + """Raised if the component is already initilalized.""" + + pass diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 4ca1a1d1f..72c35d898 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -15,7 +15,9 @@ import pkgutil import re +import icon4py.model.atmosphere.diffusion.stencils as diffusion import icon4py.model.atmosphere.dycore as dycore +import icon4py.model.common.interpolation.stencils as intp import pytest from click.testing import CliRunner @@ -25,6 +27,9 @@ DYCORE_PKG = "atmosphere.dycore" +INTERPOLATION_PKG = "common.interpolation.stencils" +DIFFUSION_PKG = "atmosphere.diffusion.stencils" + LEVELS_PER_THREAD = "1" BLOCK_SIZE = "128" OUTPATH = "." @@ -36,9 +41,21 @@ def cli(): def dycore_fencils() -> list[tuple[str, str]]: - pkgpath = os.path.dirname(dycore.__file__) + return _fencils(dycore.__file__, DYCORE_PKG) + + +def interpolation_fencils() -> list[tuple[str, str]]: + return _fencils(intp.__file__, INTERPOLATION_PKG) + + +def diffusion_fencils() -> list[tuple[str, str]]: + return _fencils(diffusion.__file__, DIFFUSION_PKG) + + +def _fencils(module_name, package_name) -> list[tuple[str, str]]: + pkgpath = os.path.dirname(module_name) stencils = [name for _, name, _ in pkgutil.iter_modules([pkgpath])] - fencils = [(DYCORE_PKG, stencil) for stencil in stencils] + fencils = [(package_name, stencil) for stencil in stencils] return fencils @@ -111,7 +128,10 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") -@pytest.mark.parametrize(("stencil_module", "stencil_name"), dycore_fencils()) +@pytest.mark.parametrize( + ("stencil_module", "stencil_name"), + dycore_fencils() + interpolation_fencils() + diffusion_fencils(), +) def test_codegen_dycore(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) with cli.isolated_filesystem(): diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 739c2ba21..55f2e5f3b 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -40,7 +40,7 @@ def test_stencil_collector_invalid_module(): def test_stencil_collector_invalid_member(): - from icon4py.model.atmosphere.dycore import apply_nabla2_to_w + from icon4py.model.atmosphere.diffusion.stencils import apply_nabla2_to_w module_path = Path(apply_nabla2_to_w.__file__) parents = module_path.parents[0] diff --git a/tools/tests/py2f/test_cffi_utils.py b/tools/tests/py2f/test_cffi_utils.py new file mode 100644 index 000000000..d8fdf6196 --- /dev/null +++ b/tools/tests/py2f/test_cffi_utils.py @@ -0,0 +1,80 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import cffi +import numpy as np +import pytest +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import float32, float64, int32, int64 +from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim, VertexDim + +from icon4pytools.py2f.cffi_utils import UnknownDimensionException, to_fields + + +n_vertices = 9 +n_edges = 27 +levels = 12 +e2c_sparse_size = 2 + + +def random(*sizes): + return np.random.default_rng().uniform(size=sizes) + + +@pytest.mark.parametrize("pointer_type", ["float*", "double*"]) +def test_unpack_from_buffer_to_field(pointer_type: str): + @to_fields(dim_sizes={VertexDim: n_vertices, KDim: levels}) + def identity( + f1: Field[[VertexDim, KDim], float], factor: float + ) -> tuple[float, Field[[VertexDim, KDim], float]]: + return factor, f1 + + ffi = cffi.FFI() + input_array = random(n_vertices, levels) + input_factor = 0.5 + res_factor, result_field = identity(ffi.from_buffer(pointer_type, input_array), input_factor) + assert res_factor == input_factor + assert np.allclose(np.asarray(result_field), input_array) + + +def test_unpack_only_scalar_args(): + @to_fields(dim_sizes={}) + def multiply(f1: float, f2: float32, f3: float64, i3: int64, i2: int32, i1: int): + return f1 * f2 * f3, i1 * i2 * i3 + + f_res, i_res = multiply(0.5, 2.0, 3.0, 1, 2, 3) + assert f_res == 3.0 + assert i_res == 6 + + +@pytest.mark.parametrize("field_type", [int32, int, int64]) +def test_unpack_local_field(field_type): + ffi = cffi.FFI() + + @to_fields(dim_sizes={EdgeDim: n_edges, E2CDim: e2c_sparse_size}) + def local_field(f1: Field[[EdgeDim, E2CDim], field_type]): + return f1 + + input_field = np.arange(n_edges * e2c_sparse_size).reshape((n_edges, e2c_sparse_size)) + res_field = local_field(ffi.from_buffer("int*", input_field)) + assert np.all(res_field == input_field) + + +def test_unknown_dimension_raises_exception(): + @to_fields(dim_sizes={}) + def do_nothing(f1: Field[[VertexDim], float]): + pass + + input_array = random() + with pytest.raises(UnknownDimensionException, match=r"size of dimension "): + do_nothing(input_array) diff --git a/tools/tests/py2f/test_code_generation.py b/tools/tests/py2f/test_code_generation.py new file mode 100644 index 000000000..74a9326de --- /dev/null +++ b/tools/tests/py2f/test_code_generation.py @@ -0,0 +1,140 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import string + +import pytest +from gt4py.next.type_system.type_specifications import ScalarKind + +from icon4pytools.py2f.codegen import ( + CffiPlugin, + CHeaderGenerator, + DimensionType, + F90InterfaceGenerator, + Func, + FuncParameter, + as_f90_value, + as_field, +) + + +field_2d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13), DimensionType(name="J", length=13)], +) +field_1d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13)], +) + +simple_type = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) + + +@pytest.mark.parametrize( + ("param", "expected"), ((simple_type, "value, "), (field_2d, ""), (field_1d, "")) +) +def test_as_target(param, expected): + assert expected == as_f90_value(param) + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:,:)"))) +def test_field_extension_2d(lang, expected): + assert as_field(field_2d, lang) == expected + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:)"))) +def test_field_extension_1d(lang, expected): + assert as_field(field_1d, lang) == expected + + +@pytest.mark.parametrize("lang", ("C", "F")) +def test_is_field_simple_type(lang): + assert as_field(simple_type, lang) == "" + + +foo = Func( + name="foo", + args=[ + FuncParameter(name="one", d_type=ScalarKind.INT32, dimensions=[]), + FuncParameter(name="two", d_type=ScalarKind.FLOAT64, dimensions=[]), + ], +) + +bar = Func( + name="bar", + args=[ + FuncParameter( + name="one", + d_type=ScalarKind.FLOAT32, + dimensions=[ + DimensionType(name="KDim", length=10), + DimensionType(name="VDim", length=50000), + ], + ), + FuncParameter(name="two", d_type=ScalarKind.INT32, dimensions=[]), + ], +) + + +def test_cheader_generation_for_single_function(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void foo(int one, double two);\n" + + +def test_cheader_for_pointer_args(): + functions = [bar] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void bar(float* one, int two);\n" + + +def compare_ignore_whitespace(s1: str, s2: str): + no_whitespace = {ord(c): None for c in string.whitespace} + return s1.translate(no_whitespace) == s2.translate(no_whitespace) + + +def test_c_header_with_several_functions(): + functions = [bar, foo] + plugin = CffiPlugin(name="libtest", functions=functions) + header = CHeaderGenerator.apply(plugin) + assert ( + header + == """extern void bar(float* one, int two);\nextern void foo(int one, double two);\n""" + ) + + +def test_fortran_interface(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + interface = F90InterfaceGenerator.apply(plugin) + expected = """ + module libtest + use, intrinsic:: iso_c_binding + implicit none + + public + interface + subroutine foo(one, two) bind(c, name='foo') + use iso_c_binding + integer(c_int), value, intent(inout):: one + real(c_double), value, intent(inout):: two + end subroutine foo + end interface + end module + """ + assert compare_ignore_whitespace(interface, expected) diff --git a/tools/tests/py2f/test_parsing_wrapper.py b/tools/tests/py2f/test_parsing_wrapper.py new file mode 100644 index 000000000..ea1952038 --- /dev/null +++ b/tools/tests/py2f/test_parsing_wrapper.py @@ -0,0 +1,39 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.parsing import parse_functions_from_module + + +def test_parse_functions(): + path = "icon4pytools.py2f.wrappers.diffusion_wrapper" + plugin = parse_functions_from_module(path) + + assert plugin.name == "diffusion_wrapper" + assert len(plugin.functions) == 2 + assert "diffusion_init" in map(lambda f: f.name, plugin.functions) + assert "diffusion_run" in map(lambda f: f.name, plugin.functions) + + +@CffiMethod.register +def do_foo(foo: str): + return foo + + +@CffiMethod.register +def do_bar(): + return "bar" + + +def test_register_with_cffi(): + assert "do_foo" in CffiMethod.get(__name__) + assert "do_bar" in CffiMethod.get(__name__) diff --git a/tools/tests/py2f/test_py2fgen.py b/tools/tests/py2f/test_py2fgen.py new file mode 100644 index 000000000..8330e4a26 --- /dev/null +++ b/tools/tests/py2f/test_py2fgen.py @@ -0,0 +1,25 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from click.testing import CliRunner + +from icon4pytools.py2f.py2fgen import main + + +def test_py2fgen(): + cli = CliRunner() + module = "icon4pytools.py2f.wrappers.diffusion_wrapper" + build_path = "./build" + with cli.isolated_filesystem(): + result = cli.invoke(main, [module, build_path]) + assert result.exit_code == 0 diff --git a/tox.ini b/tox.ini index 65d9765d1..4b8680a57 100644 --- a/tox.ini +++ b/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/common/src tools/src + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/driver/src model/common/src tools/src pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html